diff --git a/CHANGELOG.md b/CHANGELOG.md index 17a06f85..2f9f4d08 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -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.6–4.8 models - **xAI**: Grok Imagine video generation (`/v1/videos`) + CLI - **CLI tools**: Grok Build setup — writes `[model.9router]` to `~/.grok/config.toml` diff --git a/open-sse/config/userTokenLimits.js b/open-sse/config/userTokenLimits.js new file mode 100644 index 00000000..14cd31d3 --- /dev/null +++ b/open-sse/config/userTokenLimits.js @@ -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; \ No newline at end of file diff --git a/src/app/(dashboard)/dashboard/users/page.js b/src/app/(dashboard)/dashboard/users/page.js index da038cdc..82b1cd33 100644 --- a/src/app/(dashboard)/dashboard/users/page.js +++ b/src/app/(dashboard)/dashboard/users/page.js @@ -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
Loading user management…
; } @@ -158,6 +245,7 @@ export default function UsersPage() { {formatDate(entry.createdAt)}
+ {entry.role === "user" ? : null}
@@ -183,6 +271,73 @@ export default function UsersPage() { + !limitsSaving && setLimitEditor(null)} + title={`Token limits · ${limitEditor?.username || "user"}`} + footer={<>} + > +
+
+
+ hourglass_top +
+

Total token budgets

+

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.

+
+
+
+ + {limitsError ?

{limitsError}

: null} + + {limitsLoading ? ( +
Loading token limits…
+ ) : ( +
+ {TOKEN_LIMIT_PROVIDER_OPTIONS.map((provider) => ( +
+
+
+ {provider.icon} +
+
+

{provider.name}

+

{provider.description}

+
+
+
+
+ + updateTokenLimit(provider.id, USER_TOKEN_LIMIT_WINDOWS.SESSION, event.target.value)} + aria-label={`${provider.name} session token limit`} + /> +
+
+ + updateTokenLimit(provider.id, USER_TOKEN_LIMIT_WINDOWS.WEEKLY, event.target.value)} + aria-label={`${provider.name} weekly token limit`} + /> +
+
+
+ ))} +
+ )} +
+ + !saving && setDeleteTarget(null)} diff --git a/src/app/api/users/[userId]/token-limits/route.js b/src/app/api/users/[userId]/token-limits/route.js new file mode 100644 index 00000000..a3e62fea --- /dev/null +++ b/src/app/api/users/[userId]/token-limits/route.js @@ -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); + } +} diff --git a/src/lib/db/index.js b/src/lib/db/index.js index 602d635c..ff478c18 100644 --- a/src/lib/db/index.js +++ b/src/lib/db/index.js @@ -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), ); diff --git a/src/lib/db/migrations/008-user-token-limits.js b/src/lib/db/migrations/008-user-token-limits.js new file mode 100644 index 00000000..27cb7e2f --- /dev/null +++ b/src/lib/db/migrations/008-user-token-limits.js @@ -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; \ No newline at end of file diff --git a/src/lib/db/migrations/index.js b/src/lib/db/migrations/index.js index a43c71a1..716038f5 100644 --- a/src/lib/db/migrations/index.js +++ b/src/lib/db/migrations/index.js @@ -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; diff --git a/src/lib/db/repos/userTokenLimitsRepo.js b/src/lib/db/repos/userTokenLimitsRepo.js new file mode 100644 index 00000000..e4a8c0c6 --- /dev/null +++ b/src/lib/db/repos/userTokenLimitsRepo.js @@ -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); +} \ No newline at end of file diff --git a/src/lib/db/schema.js b/src/lib/db/schema.js index c2b5b0f9..ac841f31 100644 --- a/src/lib/db/schema.js +++ b/src/lib/db/schema.js @@ -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: { diff --git a/src/lib/tokenLimitEnforcer.js b/src/lib/tokenLimitEnforcer.js new file mode 100644 index 00000000..04fe6795 --- /dev/null +++ b/src/lib/tokenLimitEnforcer.js @@ -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; +} diff --git a/src/sse/handlers/chat.js b/src/sse/handlers/chat.js index 96744c5e..63445e2b 100644 --- a/src/sse/handlers/chat.js +++ b/src/sse/handlers/chat.js @@ -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 diff --git a/tests/unit/db-migration-chain.test.js b/tests/unit/db-migration-chain.test.js index e95486d3..29e5fc05 100644 --- a/tests/unit/db-migration-chain.test.js +++ b/tests/unit/db-migration-chain.test.js @@ -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 () => { diff --git a/tests/unit/user-token-limit-enforcer.test.js b/tests/unit/user-token-limit-enforcer.test.js new file mode 100644 index 00000000..b820fa73 --- /dev/null +++ b/tests/unit/user-token-limit-enforcer.test.js @@ -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(); + }); +}); diff --git a/tests/unit/user-token-limits-repo.test.js b/tests/unit/user-token-limits-repo.test.js new file mode 100644 index 00000000..7da28bd3 --- /dev/null +++ b/tests/unit/user-token-limits-repo.test.js @@ -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"); + }); +}); diff --git a/tests/unit/user-token-limits-route.test.js b/tests/unit/user-token-limits-route.test.js new file mode 100644 index 00000000..d184bfcb --- /dev/null +++ b/tests/unit/user-token-limits-route.test.js @@ -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(); + }); +});