feat: limit quota for user

This commit is contained in:
2026-07-17 18:08:50 +07:00
parent 2e595627d0
commit b700a070c0
15 changed files with 770 additions and 3 deletions
+1
View File
@@ -1,6 +1,7 @@
# v0.5.35 (2026-07-16)
## Features
- **User limits**: add per-user total-token budgets for Orbit Provider and Codex with rolling 5-hour and weekly windows
- **Orbit Provider**: add Anthropic-compatible API-key routing for Claude Opus 4.64.8 models
- **xAI**: Grok Imagine video generation (`/v1/videos`) + CLI
- **CLI tools**: Grok Build setup — writes `[model.9router]` to `~/.grok/config.toml`
+19
View File
@@ -0,0 +1,19 @@
export const USER_TOKEN_LIMIT_PROVIDERS = Object.freeze({
ORBIT: "orbit-provider",
CODEX: "codex",
});
export const USER_TOKEN_LIMIT_WINDOWS = Object.freeze({
SESSION: "session",
WEEKLY: "weekly",
});
export const USER_TOKEN_LIMIT_PROVIDER_IDS = Object.freeze(
Object.values(USER_TOKEN_LIMIT_PROVIDERS),
);
export const USER_TOKEN_LIMIT_WINDOW_IDS = Object.freeze(
Object.values(USER_TOKEN_LIMIT_WINDOWS),
);
export const USER_TOKEN_LIMIT_SESSION_MS = 5 * 60 * 60 * 1000;
+155
View File
@@ -6,8 +6,36 @@ import { Button, Card, Input } from "@/shared/components";
import Modal, { ConfirmModal } from "@/shared/components/Modal";
import useUserStore from "@/store/userStore";
import { formatVietnamDateTime } from "@/shared/utils/dateTime";
import {
USER_TOKEN_LIMIT_PROVIDERS,
USER_TOKEN_LIMIT_WINDOWS,
} from "open-sse/config/userTokenLimits.js";
const EMPTY_FORM = { username: "", password: "", role: "user", isActive: true };
const TOKEN_LIMIT_PROVIDER_OPTIONS = [
{
id: USER_TOKEN_LIMIT_PROVIDERS.ORBIT,
name: "Orbit Provider",
description: "Anthropic-compatible traffic routed through Orbit.",
icon: "orbit",
},
{
id: USER_TOKEN_LIMIT_PROVIDERS.CODEX,
name: "Codex",
description: "OpenAI Codex responses and coding sessions.",
icon: "terminal",
},
];
function createEmptyTokenLimits() {
return Object.fromEntries(TOKEN_LIMIT_PROVIDER_OPTIONS.map(({ id }) => [
id,
{
[USER_TOKEN_LIMIT_WINDOWS.SESSION]: 0,
[USER_TOKEN_LIMIT_WINDOWS.WEEKLY]: 0,
},
]));
}
function formatDate(value) {
if (!value) return "—";
@@ -25,6 +53,11 @@ export default function UsersPage() {
const [editor, setEditor] = useState(null);
const [form, setForm] = useState(EMPTY_FORM);
const [deleteTarget, setDeleteTarget] = useState(null);
const [limitEditor, setLimitEditor] = useState(null);
const [tokenLimits, setTokenLimits] = useState(createEmptyTokenLimits);
const [limitsLoading, setLimitsLoading] = useState(false);
const [limitsSaving, setLimitsSaving] = useState(false);
const [limitsError, setLimitsError] = useState("");
const loadUsers = useCallback(async () => {
setLoading(true);
@@ -113,6 +146,60 @@ export default function UsersPage() {
}
};
const openTokenLimits = async (target) => {
setLimitEditor(target);
setTokenLimits(createEmptyTokenLimits());
setLimitsError("");
setLimitsLoading(true);
try {
const response = await fetch(`/api/users/${target.id}/token-limits`, { cache: "no-store" });
const data = await response.json();
if (!response.ok) throw new Error(data.error || "Failed to load token limits");
setTokenLimits(data.limits || createEmptyTokenLimits());
} catch (requestError) {
setLimitsError(requestError.message || "Failed to load token limits");
} finally {
setLimitsLoading(false);
}
};
const updateTokenLimit = (provider, windowType, value) => {
setTokenLimits((current) => ({
...current,
[provider]: {
...current[provider],
[windowType]: value,
},
}));
};
const saveTokenLimits = async () => {
if (!limitEditor) return;
setLimitsSaving(true);
setLimitsError("");
try {
const normalizedLimits = Object.fromEntries(TOKEN_LIMIT_PROVIDER_OPTIONS.map(({ id }) => [
id,
{
[USER_TOKEN_LIMIT_WINDOWS.SESSION]: Number(tokenLimits[id]?.session || 0),
[USER_TOKEN_LIMIT_WINDOWS.WEEKLY]: Number(tokenLimits[id]?.weekly || 0),
},
]));
const response = await fetch(`/api/users/${limitEditor.id}/token-limits`, {
method: "PUT",
headers: { "Content-Type": "application/json" },
body: JSON.stringify({ limits: normalizedLimits }),
});
const data = await response.json();
if (!response.ok) throw new Error(data.error || "Failed to save token limits");
setLimitEditor(null);
} catch (requestError) {
setLimitsError(requestError.message || "Failed to save token limits");
} finally {
setLimitsSaving(false);
}
};
if (!user || user.role !== "admin") {
return <div className="py-12 text-center text-text-muted">Loading user management</div>;
}
@@ -158,6 +245,7 @@ export default function UsersPage() {
<td className="px-5 py-4 text-text-muted">{formatDate(entry.createdAt)}</td>
<td className="px-5 py-4 text-right">
<div className="flex justify-end gap-2">
{entry.role === "user" ? <Button variant="ghost" size="sm" onClick={() => openTokenLimits(entry)}>Token limits</Button> : null}
<Button variant="ghost" size="sm" onClick={() => openEdit(entry)}>Edit</Button>
<Button variant="ghost" size="sm" className="text-red-600 hover:text-red-700" onClick={() => setDeleteTarget(entry)} disabled={entry.id === user.id}>Delete</Button>
</div>
@@ -183,6 +271,73 @@ export default function UsersPage() {
</form>
</Modal>
<Modal
isOpen={!!limitEditor}
onClose={() => !limitsSaving && setLimitEditor(null)}
title={`Token limits · ${limitEditor?.username || "user"}`}
footer={<><Button variant="ghost" onClick={() => setLimitEditor(null)} disabled={limitsSaving}>Cancel</Button><Button variant="primary" onClick={saveTokenLimits} loading={limitsSaving} disabled={limitsLoading}>Save limits</Button></>}
>
<div className="space-y-5">
<div className="rounded-xl border border-brand-500/20 bg-brand-500/5 px-4 py-3">
<div className="flex items-start gap-3">
<span className="material-symbols-outlined mt-0.5 text-[20px] text-brand-500">hourglass_top</span>
<div>
<p className="text-sm font-medium text-text-main">Total token budgets</p>
<p className="mt-1 text-xs leading-5 text-text-muted">Usage includes input and output tokens. Session usage is measured over the previous 5 hours; weekly usage resets Monday at 00:00 Vietnam time. Enter 0 for unlimited.</p>
</div>
</div>
</div>
{limitsError ? <p className="rounded-lg border border-red-500/30 bg-red-500/10 px-3 py-2 text-sm text-red-600 dark:text-red-400">{limitsError}</p> : null}
{limitsLoading ? (
<div className="py-10 text-center text-sm text-text-muted">Loading token limits</div>
) : (
<div className="space-y-3">
{TOKEN_LIMIT_PROVIDER_OPTIONS.map((provider) => (
<section key={provider.id} className="rounded-xl border border-border-subtle bg-surface-2/35 p-4">
<div className="mb-4 flex items-center gap-3">
<div className="flex size-10 items-center justify-center rounded-lg border border-border-subtle bg-surface text-brand-500 shadow-sm">
<span className="material-symbols-outlined text-[21px]">{provider.icon}</span>
</div>
<div>
<h2 className="text-sm font-semibold text-text-main">{provider.name}</h2>
<p className="text-xs text-text-muted">{provider.description}</p>
</div>
</div>
<div className="grid gap-3 sm:grid-cols-2">
<div className="space-y-1.5">
<label className="text-xs font-medium uppercase tracking-wide text-text-muted">Session · 5 hours</label>
<Input
type="number"
min="0"
step="1"
inputMode="numeric"
value={tokenLimits[provider.id]?.[USER_TOKEN_LIMIT_WINDOWS.SESSION] ?? 0}
onChange={(event) => updateTokenLimit(provider.id, USER_TOKEN_LIMIT_WINDOWS.SESSION, event.target.value)}
aria-label={`${provider.name} session token limit`}
/>
</div>
<div className="space-y-1.5">
<label className="text-xs font-medium uppercase tracking-wide text-text-muted">Weekly</label>
<Input
type="number"
min="0"
step="1"
inputMode="numeric"
value={tokenLimits[provider.id]?.[USER_TOKEN_LIMIT_WINDOWS.WEEKLY] ?? 0}
onChange={(event) => updateTokenLimit(provider.id, USER_TOKEN_LIMIT_WINDOWS.WEEKLY, event.target.value)}
aria-label={`${provider.name} weekly token limit`}
/>
</div>
</div>
</section>
))}
</div>
)}
</div>
</Modal>
<ConfirmModal
isOpen={!!deleteTarget}
onClose={() => !saving && setDeleteTarget(null)}
@@ -0,0 +1,56 @@
import { NextResponse } from "next/server";
import {
getUserById,
getUserTokenLimits,
replaceUserTokenLimits,
} from "@/lib/db/index.js";
import { requireAdminUser } from "@/lib/auth/currentUser.js";
const NO_STORE_HEADERS = { "Cache-Control": "no-store" };
function errorResponse(error) {
const message = error?.message || "Request failed";
const status = message === "Unauthorized"
? 401
: message === "Forbidden"
? 403
: message === "User not found"
? 404
: 400;
return NextResponse.json({ error: message }, { status, headers: NO_STORE_HEADERS });
}
async function getTargetUser(params) {
const { userId } = await params;
const user = await getUserById(userId);
if (!user) throw new Error("User not found");
return user;
}
export async function GET(_request, { params }) {
try {
await requireAdminUser();
const user = await getTargetUser(params);
const limits = await getUserTokenLimits(user.id);
return NextResponse.json({ limits }, { headers: NO_STORE_HEADERS });
} catch (error) {
return errorResponse(error);
}
}
export async function PUT(request, { params }) {
try {
await requireAdminUser();
const user = await getTargetUser(params);
if (user.role !== "user") throw new Error("Token limits only apply to user accounts");
const body = await request.json();
if (!body?.limits || typeof body.limits !== "object" || Array.isArray(body.limits)) {
throw new Error("Token limits are required");
}
const limits = await replaceUserTokenLimits(user.id, body?.limits);
return NextResponse.json({ limits }, { headers: NO_STORE_HEADERS });
} catch (error) {
return errorResponse(error);
}
}
+28
View File
@@ -2,6 +2,13 @@
import { getAdapter } from "./driver.js";
import { stringifyJson, parseJson } from "./helpers/jsonCol.js";
import { normalizeCliToolConfig, isPersistableCliTool } from "@/shared/constants/cliToolConfig.js";
import {
USER_TOKEN_LIMIT_PROVIDER_IDS,
USER_TOKEN_LIMIT_WINDOW_IDS,
} from "open-sse/config/userTokenLimits.js";
const userTokenLimitProviderSet = new Set(USER_TOKEN_LIMIT_PROVIDER_IDS);
const userTokenLimitWindowSet = new Set(USER_TOKEN_LIMIT_WINDOW_IDS);
// Settings
export {
@@ -52,6 +59,12 @@ export {
upsertCliToolConfig, deleteCliToolConfigsByOwnerId,
} from "./repos/cliToolConfigsRepo.js";
// Per-user provider token limits
export {
createEmptyUserTokenLimits, getUserTokenLimits,
replaceUserTokenLimits, getUserProviderTokenUsageSince,
} from "./repos/userTokenLimitsRepo.js";
// Aliases (model + custom + mitm)
export {
getModelAliases, setModelAlias, deleteModelAlias,
@@ -95,6 +108,7 @@ export async function exportDb() {
apiKeys: db.all(`SELECT * FROM apiKeys`).map((r) => ({ id: r.id, key: r.key, name: r.name, machineId: r.machineId, ownerId: r.ownerId, isActive: r.isActive === 1, createdAt: r.createdAt })),
combos: db.all(`SELECT * FROM combos`).map((r) => ({ id: r.id, name: r.name, ownerId: r.ownerId, kind: r.kind, models: parseJson(r.models, []), createdAt: r.createdAt, updatedAt: r.updatedAt })),
cliToolConfigs: db.all(`SELECT * FROM cliToolConfigs`).map((r) => ({ ownerId: r.ownerId, toolId: r.toolId, config: parseJson(r.data, {}), createdAt: r.createdAt, updatedAt: r.updatedAt })),
userTokenLimits: db.all(`SELECT userId, provider, windowType, tokenLimit, createdAt, updatedAt FROM userTokenLimits`),
modelAliases: {},
customModels: [],
mitmAlias: {},
@@ -130,6 +144,7 @@ export async function importDb(payload) {
// Wipe all tables (keep _meta)
db.run(`DELETE FROM settings`);
db.run(`DELETE FROM cliToolConfigs`);
db.run(`DELETE FROM userTokenLimits`);
// Old backups predate multi-user authentication. Preserve the local
// administrator unless the payload explicitly carries a users array.
if (Array.isArray(payload.users)) db.run(`DELETE FROM users`);
@@ -155,6 +170,19 @@ export async function importDb(payload) {
}
}
const importedUserIds = new Set(db.all(`SELECT id FROM users`).map((user) => user.id));
for (const limit of payload.userTokenLimits || []) {
if (!importedUserIds.has(limit?.userId)) continue;
if (!userTokenLimitProviderSet.has(limit.provider) || !userTokenLimitWindowSet.has(limit.windowType)) continue;
const tokenLimit = Number(limit.tokenLimit);
if (!Number.isSafeInteger(tokenLimit) || tokenLimit <= 0) continue;
db.run(
`INSERT OR REPLACE INTO userTokenLimits(userId, provider, windowType, tokenLimit, createdAt, updatedAt)
VALUES(?, ?, ?, ?, ?, ?)`,
[limit.userId, limit.provider, limit.windowType, tokenLimit, limit.createdAt || new Date().toISOString(), limit.updatedAt || new Date().toISOString()],
);
}
const adminOwnerIds = new Set(
db.all(`SELECT id FROM users WHERE role = 'admin'`).map((user) => user.id),
);
@@ -0,0 +1,16 @@
import { TABLES, buildCreateTableSql } from "../schema.js";
const userTokenLimitsMigration = {
version: 8,
name: "user-token-limits",
up(db) {
const definition = TABLES.userTokenLimits;
db.exec(buildCreateTableSql("userTokenLimits", definition));
for (const index of definition.indexes || []) db.exec(index);
db.exec(
"CREATE INDEX IF NOT EXISTS idx_uh_user_provider_ts ON usageHistory(userId, provider, timestamp)",
);
},
};
export default userTokenLimitsMigration;
+2 -1
View File
@@ -8,8 +8,9 @@ import m004 from "./004-provider-connection-owners.js";
import m005 from "./005-usage-user-attribution.js";
import m006 from "./006-combo-owners.js";
import m007 from "./007-admin-provider-connections.js";
import m008 from "./008-user-token-limits.js";
export const MIGRATIONS = [m001, m002, m003, m004, m005, m006, m007].sort((a, b) => a.version - b.version);
export const MIGRATIONS = [m001, m002, m003, m004, m005, m006, m007, m008].sort((a, b) => a.version - b.version);
export function latestVersion() {
return MIGRATIONS.length ? MIGRATIONS[MIGRATIONS.length - 1].version : 0;
+115
View File
@@ -0,0 +1,115 @@
import { getAdapter } from "../driver.js";
import {
USER_TOKEN_LIMIT_PROVIDER_IDS,
USER_TOKEN_LIMIT_WINDOW_IDS,
} from "open-sse/config/userTokenLimits.js";
const providerSet = new Set(USER_TOKEN_LIMIT_PROVIDER_IDS);
const windowSet = new Set(USER_TOKEN_LIMIT_WINDOW_IDS);
export function createEmptyUserTokenLimits() {
return Object.fromEntries(
USER_TOKEN_LIMIT_PROVIDER_IDS.map((provider) => [
provider,
Object.fromEntries(USER_TOKEN_LIMIT_WINDOW_IDS.map((windowType) => [windowType, 0])),
]),
);
}
function assertProvider(provider) {
if (!providerSet.has(provider)) throw new Error("Unsupported token limit provider");
}
function assertWindowType(windowType) {
if (!windowSet.has(windowType)) throw new Error("Unsupported token limit window");
}
function normalizeTokenLimit(value) {
const tokenLimit = Number(value);
if (!Number.isSafeInteger(tokenLimit) || tokenLimit < 0) {
throw new Error("Token limit must be a non-negative integer");
}
return tokenLimit;
}
function normalizeLimits(limits) {
if (!limits || typeof limits !== "object" || Array.isArray(limits)) {
throw new Error("Token limits are required");
}
for (const provider of Object.keys(limits)) assertProvider(provider);
for (const providerLimits of Object.values(limits)) {
if (!providerLimits || typeof providerLimits !== "object" || Array.isArray(providerLimits)) {
throw new Error("Provider token limits must be an object");
}
for (const windowType of Object.keys(providerLimits)) assertWindowType(windowType);
}
const normalized = createEmptyUserTokenLimits();
for (const provider of USER_TOKEN_LIMIT_PROVIDER_IDS) {
for (const windowType of USER_TOKEN_LIMIT_WINDOW_IDS) {
normalized[provider][windowType] = normalizeTokenLimit(
limits?.[provider]?.[windowType] ?? 0,
);
}
}
return normalized;
}
export async function getUserTokenLimits(userId) {
if (!userId) return createEmptyUserTokenLimits();
const db = await getAdapter();
const limits = createEmptyUserTokenLimits();
const rows = db.all(
`SELECT provider, windowType, tokenLimit
FROM userTokenLimits
WHERE userId = ?`,
[userId],
);
for (const row of rows) {
if (!providerSet.has(row.provider) || !windowSet.has(row.windowType)) continue;
limits[row.provider][row.windowType] = Math.max(0, Number(row.tokenLimit) || 0);
}
return limits;
}
export async function replaceUserTokenLimits(userId, limits) {
if (!userId) throw new Error("User id is required");
const normalized = normalizeLimits(limits);
const db = await getAdapter();
const now = new Date().toISOString();
db.transaction(() => {
db.run(`DELETE FROM userTokenLimits WHERE userId = ?`, [userId]);
for (const provider of USER_TOKEN_LIMIT_PROVIDER_IDS) {
for (const windowType of USER_TOKEN_LIMIT_WINDOW_IDS) {
const tokenLimit = normalized[provider][windowType];
if (tokenLimit === 0) continue;
db.run(
`INSERT INTO userTokenLimits(userId, provider, windowType, tokenLimit, createdAt, updatedAt)
VALUES(?, ?, ?, ?, ?, ?)`,
[userId, provider, windowType, tokenLimit, now, now],
);
}
}
});
return normalized;
}
export async function getUserProviderTokenUsageSince(userId, provider, since) {
if (!userId) return 0;
assertProvider(provider);
if (!(since instanceof Date) || !Number.isFinite(since.getTime())) {
throw new Error("A valid usage window start is required");
}
const db = await getAdapter();
const row = db.get(
`SELECT COALESCE(SUM(COALESCE(promptTokens, 0) + COALESCE(completionTokens, 0)), 0) AS totalTokens
FROM usageHistory
WHERE userId = ? AND provider = ? AND timestamp >= ?`,
[userId, provider, since.toISOString()],
);
return Math.max(0, Number(row?.totalTokens) || 0);
}
+16 -1
View File
@@ -3,7 +3,7 @@
// pre-change safety backup in migrate.js: when the stored version is lower,
// one lightweight DB backup is taken before applying schema changes. Forgetting
// to bump only skips that backup — it does NOT break the additive auto-sync.
export const SCHEMA_VERSION = 8;
export const SCHEMA_VERSION = 9;
export const PRAGMA_SQL = `
PRAGMA journal_mode = WAL;
@@ -134,6 +134,20 @@ export const TABLES = {
primaryKey: "PRIMARY KEY (ownerId, toolId)",
indexes: ["CREATE INDEX IF NOT EXISTS idx_cli_tool_configs_owner ON cliToolConfigs(ownerId)"],
},
userTokenLimits: {
columns: {
userId: "TEXT NOT NULL REFERENCES users(id) ON DELETE CASCADE",
provider: "TEXT NOT NULL",
windowType: "TEXT NOT NULL",
tokenLimit: "INTEGER NOT NULL DEFAULT 0",
createdAt: "TEXT NOT NULL",
updatedAt: "TEXT NOT NULL",
},
primaryKey: "PRIMARY KEY (userId, provider, windowType)",
indexes: [
"CREATE INDEX IF NOT EXISTS idx_user_token_limits_user ON userTokenLimits(userId)",
],
},
kv: {
columns: {
scope: "TEXT NOT NULL",
@@ -166,6 +180,7 @@ export const TABLES = {
"CREATE INDEX IF NOT EXISTS idx_uh_model ON usageHistory(model)",
"CREATE INDEX IF NOT EXISTS idx_uh_conn ON usageHistory(connectionId)",
"CREATE INDEX IF NOT EXISTS idx_uh_user ON usageHistory(userId)",
"CREATE INDEX IF NOT EXISTS idx_uh_user_provider_ts ON usageHistory(userId, provider, timestamp)",
],
},
usageDaily: {
+74
View File
@@ -0,0 +1,74 @@
import {
USER_TOKEN_LIMIT_PROVIDER_IDS,
USER_TOKEN_LIMIT_SESSION_MS,
USER_TOKEN_LIMIT_WINDOWS,
} from "open-sse/config/userTokenLimits.js";
import {
getUserById,
getUserProviderTokenUsageSince,
getUserTokenLimits,
} from "@/lib/db/index.js";
import {
getVietnamDateKey,
shiftVietnamDateKey,
} from "@/shared/utils/dateTime.js";
const limitedProviderSet = new Set(USER_TOKEN_LIMIT_PROVIDER_IDS);
export function getUserTokenLimitWindowStart(windowType, now = new Date()) {
const current = now instanceof Date ? now : new Date(now);
if (!Number.isFinite(current.getTime())) throw new Error("A valid current time is required");
if (windowType === USER_TOKEN_LIMIT_WINDOWS.SESSION) {
return new Date(current.getTime() - USER_TOKEN_LIMIT_SESSION_MS);
}
if (windowType === USER_TOKEN_LIMIT_WINDOWS.WEEKLY) {
const dateKey = getVietnamDateKey(current);
const vietnamNoon = new Date(`${dateKey}T12:00:00+07:00`);
const daysSinceMonday = (vietnamNoon.getUTCDay() + 6) % 7;
const mondayKey = shiftVietnamDateKey(dateKey, -daysSinceMonday);
return new Date(`${mondayKey}T00:00:00+07:00`);
}
throw new Error("Unsupported token limit window");
}
/**
* Check whether a dashboard user has exhausted a provider token budget.
* Returns null when the provider/user is exempt or all configured limits have headroom.
*/
export async function checkUserTokenLimit(userId, provider, now = new Date()) {
if (!userId || !limitedProviderSet.has(provider)) return null;
const user = await getUserById(userId);
if (!user || !user.isActive || user.role !== "user") return null;
const limits = await getUserTokenLimits(user.id);
const providerLimits = limits[provider];
if (!providerLimits) return null;
for (const windowType of [
USER_TOKEN_LIMIT_WINDOWS.SESSION,
USER_TOKEN_LIMIT_WINDOWS.WEEKLY,
]) {
const limit = providerLimits[windowType];
if (!Number.isSafeInteger(limit) || limit <= 0) continue;
const windowStart = getUserTokenLimitWindowStart(windowType, now);
const used = await getUserProviderTokenUsageSince(user.id, provider, windowStart);
if (used >= limit) {
return {
exceeded: true,
provider,
windowType,
limit,
used,
remaining: 0,
windowStart: windowStart.toISOString(),
};
}
}
return null;
}
+9
View File
@@ -26,6 +26,7 @@ import * as log from "../utils/logger.js";
import { updateProviderCredentials, checkAndRefreshToken } from "../services/tokenRefresh.js";
import { getProjectIdForConnection } from "open-sse/services/projectId.js";
import { getDisabledModelResponse } from "../services/disabledModels.js";
import { checkUserTokenLimit } from "@/lib/tokenLimitEnforcer.js";
/**
* Handle chat completion request
@@ -203,6 +204,14 @@ async function handleSingleModelChat(body, modelStr, clientRawRequest = null, re
const disabledModelResponse = await getDisabledModelResponse(provider, model);
if (disabledModelResponse) return disabledModelResponse;
const tokenLimitResult = await checkUserTokenLimit(ownerId, provider);
if (tokenLimitResult) {
const { windowType, limit, used } = tokenLimitResult;
const message = `Token limit exceeded: ${windowType} limit of ${limit} tokens reached (${used} used)`;
log.warn("TOKEN_LIMIT", message, { userId: ownerId, provider, windowType, limit, used });
return errorResponse(HTTP_STATUS.RATE_LIMITED, message);
}
// Routing shown in the unified "▶" line (client model → provider/model)
// Extract userAgent from request
+3 -1
View File
@@ -34,13 +34,15 @@ describe("Schema migrations", () => {
const tables = db.all(`SELECT name FROM sqlite_master WHERE type='table'`).map(t => t.name);
expect(tables).toEqual(expect.arrayContaining([
"_meta", "settings", "providerConnections", "providerNodes",
"_meta", "settings", "users", "userTokenLimits", "providerConnections", "providerNodes",
"proxyPools", "apiKeys", "combos", "kv", "usageHistory", "usageDaily", "requestDetails",
]));
expect(db.all(`PRAGMA table_info(providerConnections)`).map((column) => column.name)).toContain("ownerId");
expect(db.all(`PRAGMA index_list(providerConnections)`).map((index) => index.name)).toContain("idx_pc_owner");
expect(db.all(`PRAGMA table_info(combos)`).map((column) => column.name)).toContain("ownerId");
expect(db.all(`PRAGMA index_list(combos)`).map((index) => index.name)).toContain("idx_combo_owner_name");
expect(db.all(`PRAGMA index_list(userTokenLimits)`).map((index) => index.name)).toContain("idx_user_token_limits_user");
expect(db.all(`PRAGMA index_list(usageHistory)`).map((index) => index.name)).toContain("idx_uh_user_provider_ts");
});
it("existing DB at older schemaVersion → re-applies pending migrations on restart", async () => {
@@ -0,0 +1,93 @@
import { beforeEach, describe, expect, it, vi } from "vitest";
const getUserById = vi.fn();
const getUserProviderTokenUsageSince = vi.fn();
const getUserTokenLimits = vi.fn();
vi.mock("@/lib/db/index.js", () => ({
getUserById,
getUserProviderTokenUsageSince,
getUserTokenLimits,
}));
const {
checkUserTokenLimit,
getUserTokenLimitWindowStart,
} = await import("@/lib/tokenLimitEnforcer.js");
describe("user token limit enforcement", () => {
beforeEach(() => {
getUserById.mockReset();
getUserProviderTokenUsageSince.mockReset();
getUserTokenLimits.mockReset();
getUserById.mockResolvedValue({ id: "user-1", role: "user", isActive: true });
getUserTokenLimits.mockResolvedValue({
"orbit-provider": { session: 100, weekly: 1000 },
codex: { session: 200, weekly: 2000 },
});
});
it("calculates rolling session and Monday Vietnam weekly window starts", () => {
const now = new Date("2026-07-17T10:00:00.000Z");
expect(getUserTokenLimitWindowStart("session", now).toISOString())
.toBe("2026-07-17T05:00:00.000Z");
expect(getUserTokenLimitWindowStart("weekly", now).toISOString())
.toBe("2026-07-12T17:00:00.000Z");
});
it("blocks when the rolling session total reaches its limit", async () => {
getUserProviderTokenUsageSince.mockResolvedValueOnce(100);
const result = await checkUserTokenLimit(
"user-1",
"orbit-provider",
new Date("2026-07-17T10:00:00.000Z"),
);
expect(result).toMatchObject({
exceeded: true,
provider: "orbit-provider",
windowType: "session",
limit: 100,
used: 100,
});
expect(getUserProviderTokenUsageSince).toHaveBeenCalledTimes(1);
});
it("checks weekly usage after the session window still has headroom", async () => {
getUserProviderTokenUsageSince
.mockResolvedValueOnce(80)
.mockResolvedValueOnce(1200);
const result = await checkUserTokenLimit(
"user-1",
"orbit-provider",
new Date("2026-07-17T10:00:00.000Z"),
);
expect(result).toMatchObject({ windowType: "weekly", limit: 1000, used: 1200 });
expect(getUserProviderTokenUsageSince).toHaveBeenCalledTimes(2);
});
it("exempts administrators, unknown users, and unsupported providers", async () => {
getUserById.mockResolvedValueOnce({ id: "admin-1", role: "admin", isActive: true });
await expect(checkUserTokenLimit("admin-1", "codex")).resolves.toBeNull();
getUserById.mockResolvedValueOnce(null);
await expect(checkUserTokenLimit("missing", "codex")).resolves.toBeNull();
await expect(checkUserTokenLimit("user-1", "openai")).resolves.toBeNull();
expect(getUserTokenLimits).not.toHaveBeenCalled();
});
it("treats zero limits as unlimited", async () => {
getUserTokenLimits.mockResolvedValue({
"orbit-provider": { session: 0, weekly: 0 },
codex: { session: 0, weekly: 0 },
});
await expect(checkUserTokenLimit("user-1", "codex")).resolves.toBeNull();
expect(getUserProviderTokenUsageSince).not.toHaveBeenCalled();
});
});
+101
View File
@@ -0,0 +1,101 @@
import fs from "node:fs";
import os from "node:os";
import path from "node:path";
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
let tempDir;
const originalDataDir = process.env.DATA_DIR;
beforeEach(() => {
tempDir = fs.mkdtempSync(path.join(os.tmpdir(), "9router-token-limits-"));
process.env.DATA_DIR = tempDir;
delete global._dbAdapter;
vi.resetModules();
});
afterEach(() => {
try { global._dbAdapter?.instance?.close?.(); } catch {}
delete global._dbAdapter;
fs.rmSync(tempDir, { recursive: true, force: true });
if (originalDataDir === undefined) delete process.env.DATA_DIR;
else process.env.DATA_DIR = originalDataDir;
});
describe("user token limit repository", () => {
it("replaces and normalizes all supported provider limits", async () => {
const db = await import("@/lib/db/index.js");
const user = await db.createUser({ username: "limited-user", password: "password", role: "user" });
await expect(db.getUserTokenLimits(user.id)).resolves.toEqual({
"orbit-provider": { session: 0, weekly: 0 },
codex: { session: 0, weekly: 0 },
});
await db.replaceUserTokenLimits(user.id, {
"orbit-provider": { session: 100, weekly: 1000 },
codex: { session: 0, weekly: 2000 },
});
await expect(db.getUserTokenLimits(user.id)).resolves.toEqual({
"orbit-provider": { session: 100, weekly: 1000 },
codex: { session: 0, weekly: 2000 },
});
});
it("sums prompt and completion tokens by user, provider, and timestamp", async () => {
const db = await import("@/lib/db/index.js");
const { getAdapter } = await import("@/lib/db/driver.js");
const user = await db.createUser({ username: "usage-user", password: "password", role: "user" });
const adapter = await getAdapter();
for (const [timestamp, provider, prompt, completion] of [
["2026-07-17T06:00:00.000Z", "codex", 30, 20],
["2026-07-17T09:00:00.000Z", "codex", 40, 10],
["2026-07-17T09:30:00.000Z", "orbit-provider", 500, 500],
]) {
adapter.run(
`INSERT INTO usageHistory(timestamp, provider, userId, promptTokens, completionTokens)
VALUES(?, ?, ?, ?, ?)`,
[timestamp, provider, user.id, prompt, completion],
);
}
await expect(db.getUserProviderTokenUsageSince(
user.id,
"codex",
new Date("2026-07-17T08:00:00.000Z"),
)).resolves.toBe(50);
});
it("rejects negative and non-integer limits without changing stored values", async () => {
const db = await import("@/lib/db/index.js");
const user = await db.createUser({ username: "invalid-limit", password: "password", role: "user" });
await db.replaceUserTokenLimits(user.id, {
"orbit-provider": { session: 100, weekly: 0 },
codex: { session: 0, weekly: 0 },
});
await expect(db.replaceUserTokenLimits(user.id, {
"orbit-provider": { session: -1, weekly: 0 },
codex: { session: 0, weekly: 0 },
})).rejects.toThrow("Token limit must be a non-negative integer");
await expect(db.getUserTokenLimits(user.id)).resolves.toMatchObject({
"orbit-provider": { session: 100 },
});
});
it("rejects unsupported providers and window names", async () => {
const db = await import("@/lib/db/index.js");
const user = await db.createUser({ username: "invalid-scope", password: "password", role: "user" });
await expect(db.replaceUserTokenLimits(user.id, {
openai: { session: 100 },
})).rejects.toThrow("Unsupported token limit provider");
await expect(db.replaceUserTokenLimits(user.id, {
codex: { monthly: 100 },
})).rejects.toThrow("Unsupported token limit window");
});
});
@@ -0,0 +1,82 @@
import { beforeEach, describe, expect, it, vi } from "vitest";
const requireAdminUser = vi.fn();
const getUserById = vi.fn();
const getUserTokenLimits = vi.fn();
const replaceUserTokenLimits = vi.fn();
vi.mock("next/server", () => ({
NextResponse: {
json(body, init = {}) {
return new Response(JSON.stringify(body), {
status: init.status || 200,
headers: { "Content-Type": "application/json", ...(init.headers || {}) },
});
},
},
}));
vi.mock("@/lib/auth/currentUser.js", () => ({ requireAdminUser }));
vi.mock("@/lib/db/index.js", () => ({
getUserById,
getUserTokenLimits,
replaceUserTokenLimits,
}));
const { GET, PUT } = await import("@/app/api/users/[userId]/token-limits/route.js");
const context = (userId = "user-1") => ({ params: Promise.resolve({ userId }) });
const putRequest = (body) => new Request("https://9router.local/api/users/user-1/token-limits", {
method: "PUT",
headers: { "Content-Type": "application/json" },
body: JSON.stringify(body),
});
const limits = {
"orbit-provider": { session: 100, weekly: 1000 },
codex: { session: 200, weekly: 2000 },
};
describe("/api/users/[userId]/token-limits", () => {
beforeEach(() => {
requireAdminUser.mockReset();
getUserById.mockReset();
getUserTokenLimits.mockReset();
replaceUserTokenLimits.mockReset();
requireAdminUser.mockResolvedValue({ id: "admin-1", role: "admin" });
getUserById.mockResolvedValue({ id: "user-1", role: "user", isActive: true });
getUserTokenLimits.mockResolvedValue(limits);
replaceUserTokenLimits.mockResolvedValue(limits);
});
it("returns an administrator-only no-store response", async () => {
const response = await GET(new Request("https://9router.local"), context());
expect(response.status).toBe(200);
expect(response.headers.get("Cache-Control")).toBe("no-store");
await expect(response.json()).resolves.toEqual({ limits });
expect(getUserTokenLimits).toHaveBeenCalledWith("user-1");
});
it("replaces all limits for a regular user", async () => {
const response = await PUT(putRequest({ limits }), context());
expect(response.status).toBe(200);
expect(replaceUserTokenLimits).toHaveBeenCalledWith("user-1", limits);
});
it("rejects non-admin access and administrator targets", async () => {
requireAdminUser.mockRejectedValueOnce(new Error("Forbidden"));
expect((await GET(new Request("https://9router.local"), context())).status).toBe(403);
getUserById.mockResolvedValueOnce({ id: "admin-2", role: "admin", isActive: true });
const response = await PUT(putRequest({ limits }), context("admin-2"));
expect(response.status).toBe(400);
expect(replaceUserTokenLimits).not.toHaveBeenCalled();
});
it("requires an explicit limits object", async () => {
const response = await PUT(putRequest({}), context());
expect(response.status).toBe(400);
expect(replaceUserTokenLimits).not.toHaveBeenCalled();
});
});