mirror of
https://github.com/Nezumi-2711/9router.git
synced 2026-09-22 05:31:47 +00:00
fix: logic reset quota by session
This commit is contained in:
@@ -1,5 +1,8 @@
|
||||
# v0.5.35 (2026-07-16)
|
||||
|
||||
## Fixes
|
||||
- **User quota**: reset fixed 5-hour session budgets at their scheduled boundary and keep weekly budget resets aligned to Monday 00:00 Vietnam time
|
||||
|
||||
## Features
|
||||
- **User limits**: add per-user total-token budgets for Orbit Provider and Codex with rolling 5-hour and weekly windows
|
||||
- **User quota**: show remaining Orbit and Codex headroom in the users table with session and weekly usage details
|
||||
|
||||
@@ -19,7 +19,7 @@ export const USER_TOKEN_LIMIT_WINDOW_IDS = Object.freeze(
|
||||
export const USER_TOKEN_LIMIT_WINDOW_CONFIG = Object.freeze({
|
||||
[USER_TOKEN_LIMIT_WINDOWS.SESSION]: Object.freeze({
|
||||
name: "Session",
|
||||
description: "Rolling 5 hours",
|
||||
description: "Fixed 5 hours from the first request",
|
||||
}),
|
||||
[USER_TOKEN_LIMIT_WINDOWS.WEEKLY]: Object.freeze({
|
||||
name: "Weekly",
|
||||
|
||||
@@ -36,8 +36,6 @@ function getQuotaTone(percentage) {
|
||||
return { bar: "bg-red-500", dot: "bg-red-500", text: "text-red-500" };
|
||||
}
|
||||
|
||||
const INACTIVE_SESSION_COUNTDOWN = "4h59m";
|
||||
|
||||
function TokenQuotaResetStatus({ quota }) {
|
||||
const [now, setNow] = useState(() => new Date());
|
||||
|
||||
@@ -51,7 +49,7 @@ function TokenQuotaResetStatus({ quota }) {
|
||||
const countdown = formatResetTime(quota.resetAt, now);
|
||||
const isSession = quota.windowType === USER_TOKEN_LIMIT_WINDOWS.SESSION;
|
||||
const text = countdown === "-"
|
||||
? (isSession ? INACTIVE_SESSION_COUNTDOWN : "Reset time unavailable")
|
||||
? (isSession ? "Session begins with the next request" : "Reset time unavailable")
|
||||
: (isSession ? `Next tokens restore in ${countdown}` : `Resets in ${countdown}`);
|
||||
|
||||
return (
|
||||
|
||||
@@ -358,7 +358,7 @@ export default function UsersPage() {
|
||||
<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>
|
||||
<p className="mt-1 text-xs leading-5 text-text-muted">Usage includes input and output tokens. A session lasts 5 hours from its first request; weekly usage resets Monday at 00:00 Vietnam time. Enter 0 for unlimited.</p>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
+2
-1
@@ -63,7 +63,8 @@ export {
|
||||
export {
|
||||
createEmptyUserTokenLimits, getUserTokenLimits,
|
||||
replaceUserTokenLimits, getUserProviderTokenUsageSince,
|
||||
getUserProviderEarliestTokenUsageSince,
|
||||
getUserProviderEarliestTokenUsageSince, getUserTokenQuotaSession,
|
||||
ensureUserTokenQuotaSession,
|
||||
} from "./repos/userTokenLimitsRepo.js";
|
||||
|
||||
// Aliases (model + custom + mitm)
|
||||
|
||||
@@ -0,0 +1,13 @@
|
||||
import { TABLES, buildCreateTableSql } from "../schema.js";
|
||||
|
||||
const userTokenQuotaSessionsMigration = {
|
||||
version: 10,
|
||||
name: "user-token-quota-sessions",
|
||||
up(db) {
|
||||
const definition = TABLES.userTokenQuotaSessions;
|
||||
db.exec(buildCreateTableSql("userTokenQuotaSessions", definition));
|
||||
for (const index of definition.indexes || []) db.exec(index);
|
||||
},
|
||||
};
|
||||
|
||||
export default userTokenQuotaSessionsMigration;
|
||||
@@ -10,8 +10,9 @@ import m006 from "./006-combo-owners.js";
|
||||
import m007 from "./007-admin-provider-connections.js";
|
||||
import m008 from "./008-user-token-limits.js";
|
||||
import m009 from "./009-remove-tunnel-settings.js";
|
||||
import m010 from "./010-user-token-quota-sessions.js";
|
||||
|
||||
export const MIGRATIONS = [m001, m002, m003, m004, m005, m006, m007, m008, m009].sort((a, b) => a.version - b.version);
|
||||
export const MIGRATIONS = [m001, m002, m003, m004, m005, m006, m007, m008, m009, m010].sort((a, b) => a.version - b.version);
|
||||
|
||||
export function latestVersion() {
|
||||
return MIGRATIONS.length ? MIGRATIONS[MIGRATIONS.length - 1].version : 0;
|
||||
|
||||
@@ -3,6 +3,8 @@ import { getAdapter } from "../driver.js";
|
||||
import { parseJson, stringifyJson } from "../helpers/jsonCol.js";
|
||||
import { getMeta, setMeta } from "../helpers/metaStore.js";
|
||||
import { appendUsageAccessClause, getUsageAccessScope } from "./usageAccessScope.js";
|
||||
import { ensureUserTokenQuotaSession } from "./userTokenLimitsRepo.js";
|
||||
import { USER_TOKEN_LIMIT_PROVIDER_IDS } from "open-sse/config/userTokenLimits.js";
|
||||
import {
|
||||
formatVietnamDateTime,
|
||||
getVietnamDateKey,
|
||||
@@ -20,6 +22,7 @@ const PENDING_TIMEOUT_MS = 60 * 1000;
|
||||
const RING_CAP = 50;
|
||||
const CONN_CACHE_TTL_MS = 30 * 1000;
|
||||
const PERIOD_MS = { "24h": 86400000, "7d": 604800000, "30d": 2592000000, "60d": 5184000000 };
|
||||
const userTokenQuotaProviderSet = new Set(USER_TOKEN_LIMIT_PROVIDER_IDS);
|
||||
|
||||
// In-memory state shared across Next.js modules
|
||||
if (!global._pendingRequests) global._pendingRequests = { byModel: {}, byAccount: {} };
|
||||
@@ -316,6 +319,7 @@ export async function saveRequestUsage(entry) {
|
||||
stringifyJson(tokens), stringifyJson({}),
|
||||
]
|
||||
);
|
||||
inserted = true;
|
||||
|
||||
const dateKey = getLocalDateKey(entry.timestamp);
|
||||
const row = db.get(`SELECT data FROM usageDaily WHERE dateKey = ?`, [dateKey]);
|
||||
@@ -323,6 +327,7 @@ export async function saveRequestUsage(entry) {
|
||||
requests: 0, promptTokens: 0, completionTokens: 0, cost: 0,
|
||||
byProvider: {}, byModel: {}, byAccount: {}, byApiKey: {}, byUser: {}, byEndpoint: {},
|
||||
};
|
||||
|
||||
aggregateEntryToDay(day, entry);
|
||||
db.run(`INSERT INTO usageDaily(dateKey, data) VALUES(?, ?) ON CONFLICT(dateKey) DO UPDATE SET data = excluded.data`, [dateKey, stringifyJson(day)]);
|
||||
|
||||
@@ -333,6 +338,10 @@ export async function saveRequestUsage(entry) {
|
||||
inserted = true;
|
||||
});
|
||||
|
||||
if (inserted && entry.userId && userTokenQuotaProviderSet.has(entry.provider) && (promptTokens + completionTokens) > 0) {
|
||||
await ensureUserTokenQuotaSession(entry.userId, entry.provider, entry.timestamp);
|
||||
}
|
||||
|
||||
if (inserted) {
|
||||
pushToRing(entry);
|
||||
scheduleStatsEvent("update", 250);
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
import { getAdapter } from "../driver.js";
|
||||
import {
|
||||
USER_TOKEN_LIMIT_PROVIDER_IDS,
|
||||
USER_TOKEN_LIMIT_SESSION_MS,
|
||||
USER_TOKEN_LIMIT_WINDOW_IDS,
|
||||
} from "open-sse/config/userTokenLimits.js";
|
||||
|
||||
@@ -137,3 +138,67 @@ export async function getUserProviderEarliestTokenUsageSince(userId, provider, s
|
||||
);
|
||||
return row?.timestamp || null;
|
||||
}
|
||||
|
||||
export async function getUserTokenQuotaSession(userId, provider) {
|
||||
if (!userId) return null;
|
||||
assertProvider(provider);
|
||||
|
||||
const db = await getAdapter();
|
||||
const row = db.get(
|
||||
`SELECT sessionStartedAt
|
||||
FROM userTokenQuotaSessions
|
||||
WHERE userId = ? AND provider = ?`,
|
||||
[userId, provider],
|
||||
);
|
||||
return row?.sessionStartedAt || null;
|
||||
}
|
||||
|
||||
/**
|
||||
* Start a fixed five-hour session only when no active one exists. Once its
|
||||
* reset time passes, the next billable request starts the following session.
|
||||
*/
|
||||
export async function ensureUserTokenQuotaSession(userId, provider, sessionStartedAt) {
|
||||
if (!userId) return null;
|
||||
assertProvider(provider);
|
||||
|
||||
const nextStart = new Date(sessionStartedAt);
|
||||
if (!Number.isFinite(nextStart.getTime())) {
|
||||
throw new Error("A valid session start time is required");
|
||||
}
|
||||
|
||||
const db = await getAdapter();
|
||||
const nextStartIso = nextStart.toISOString();
|
||||
let activeSessionStart = null;
|
||||
|
||||
db.transaction(() => {
|
||||
const row = db.get(
|
||||
`SELECT sessionStartedAt
|
||||
FROM userTokenQuotaSessions
|
||||
WHERE userId = ? AND provider = ?`,
|
||||
[userId, provider],
|
||||
);
|
||||
const currentStart = row?.sessionStartedAt ? new Date(row.sessionStartedAt) : null;
|
||||
const currentIsActive = Number.isFinite(currentStart?.getTime())
|
||||
&& (
|
||||
currentStart.getTime() >= nextStart.getTime()
|
||||
|| currentStart.getTime() + USER_TOKEN_LIMIT_SESSION_MS > nextStart.getTime()
|
||||
);
|
||||
|
||||
if (currentIsActive) {
|
||||
activeSessionStart = currentStart.toISOString();
|
||||
return;
|
||||
}
|
||||
|
||||
db.run(
|
||||
`INSERT INTO userTokenQuotaSessions(userId, provider, sessionStartedAt, updatedAt)
|
||||
VALUES(?, ?, ?, ?)
|
||||
ON CONFLICT(userId, provider) DO UPDATE SET
|
||||
sessionStartedAt = excluded.sessionStartedAt,
|
||||
updatedAt = excluded.updatedAt`,
|
||||
[userId, provider, nextStartIso, new Date().toISOString()],
|
||||
);
|
||||
activeSessionStart = nextStartIso;
|
||||
});
|
||||
|
||||
return activeSessionStart;
|
||||
}
|
||||
@@ -148,6 +148,18 @@ export const TABLES = {
|
||||
"CREATE INDEX IF NOT EXISTS idx_user_token_limits_user ON userTokenLimits(userId)",
|
||||
],
|
||||
},
|
||||
userTokenQuotaSessions: {
|
||||
columns: {
|
||||
userId: "TEXT NOT NULL REFERENCES users(id) ON DELETE CASCADE",
|
||||
provider: "TEXT NOT NULL",
|
||||
sessionStartedAt: "TEXT NOT NULL",
|
||||
updatedAt: "TEXT NOT NULL",
|
||||
},
|
||||
primaryKey: "PRIMARY KEY (userId, provider)",
|
||||
indexes: [
|
||||
"CREATE INDEX IF NOT EXISTS idx_user_token_quota_sessions_user ON userTokenQuotaSessions(userId)",
|
||||
],
|
||||
},
|
||||
kv: {
|
||||
columns: {
|
||||
scope: "TEXT NOT NULL",
|
||||
|
||||
@@ -1,17 +1,20 @@
|
||||
import {
|
||||
USER_TOKEN_LIMIT_PROVIDER_IDS,
|
||||
USER_TOKEN_LIMIT_SESSION_MS,
|
||||
USER_TOKEN_LIMIT_WINDOWS,
|
||||
} from "open-sse/config/userTokenLimits.js";
|
||||
import {
|
||||
ensureUserTokenQuotaSession,
|
||||
getUserById,
|
||||
getUserProviderEarliestTokenUsageSince,
|
||||
getUserProviderTokenUsageSince,
|
||||
getUserTokenQuotaSession,
|
||||
getUserTokenLimits,
|
||||
} from "@/lib/db/index.js";
|
||||
import {
|
||||
getVietnamDateKey,
|
||||
shiftVietnamDateKey,
|
||||
} from "@/shared/utils/dateTime.js";
|
||||
getActiveSessionWindowStart,
|
||||
getRollingSessionWindowStart,
|
||||
getWeeklyTokenLimitWindowStart,
|
||||
} from "@/lib/userTokenLimitWindows.js";
|
||||
|
||||
const limitedProviderSet = new Set(USER_TOKEN_LIMIT_PROVIDER_IDS);
|
||||
|
||||
@@ -20,20 +23,39 @@ export function getUserTokenLimitWindowStart(windowType, now = new Date()) {
|
||||
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);
|
||||
return getRollingSessionWindowStart(current);
|
||||
}
|
||||
|
||||
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`);
|
||||
return getWeeklyTokenLimitWindowStart(current);
|
||||
}
|
||||
|
||||
throw new Error("Unsupported token limit window");
|
||||
}
|
||||
|
||||
async function getActiveSessionStart(userId, provider, now) {
|
||||
const storedSessionStart = await getUserTokenQuotaSession(userId, provider);
|
||||
if (storedSessionStart) {
|
||||
return getActiveSessionWindowStart(storedSessionStart, now);
|
||||
}
|
||||
|
||||
// Keep active usage for installations created before fixed sessions existed.
|
||||
const earliestUsageAt = await getUserProviderEarliestTokenUsageSince(
|
||||
userId,
|
||||
provider,
|
||||
getRollingSessionWindowStart(now),
|
||||
);
|
||||
const legacySessionStart = getActiveSessionWindowStart(earliestUsageAt, now);
|
||||
if (!legacySessionStart) return null;
|
||||
|
||||
const savedSessionStart = await ensureUserTokenQuotaSession(
|
||||
userId,
|
||||
provider,
|
||||
legacySessionStart,
|
||||
);
|
||||
return getActiveSessionWindowStart(savedSessionStart, now);
|
||||
}
|
||||
|
||||
/**
|
||||
* 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.
|
||||
@@ -55,8 +77,12 @@ export async function checkUserTokenLimit(userId, provider, now = new Date()) {
|
||||
const limit = providerLimits[windowType];
|
||||
if (!Number.isSafeInteger(limit) || limit <= 0) continue;
|
||||
|
||||
const windowStart = getUserTokenLimitWindowStart(windowType, now);
|
||||
const used = await getUserProviderTokenUsageSince(user.id, provider, windowStart);
|
||||
const windowStart = windowType === USER_TOKEN_LIMIT_WINDOWS.SESSION
|
||||
? await getActiveSessionStart(user.id, provider, now)
|
||||
: getWeeklyTokenLimitWindowStart(now);
|
||||
const used = windowStart
|
||||
? await getUserProviderTokenUsageSince(user.id, provider, windowStart)
|
||||
: 0;
|
||||
if (used >= limit) {
|
||||
return {
|
||||
exceeded: true,
|
||||
|
||||
@@ -0,0 +1,44 @@
|
||||
import {
|
||||
USER_TOKEN_LIMIT_SESSION_MS,
|
||||
} from "open-sse/config/userTokenLimits.js";
|
||||
import {
|
||||
getVietnamDateKey,
|
||||
shiftVietnamDateKey,
|
||||
} from "@/shared/utils/dateTime.js";
|
||||
|
||||
function toValidDate(value) {
|
||||
if (value === null || value === undefined || value === "") return null;
|
||||
const date = value instanceof Date ? value : new Date(value);
|
||||
return Number.isFinite(date.getTime()) ? date : null;
|
||||
}
|
||||
|
||||
export function getRollingSessionWindowStart(now = new Date()) {
|
||||
const current = toValidDate(now);
|
||||
if (!current) throw new Error("A valid current time is required");
|
||||
return new Date(current.getTime() - USER_TOKEN_LIMIT_SESSION_MS);
|
||||
}
|
||||
|
||||
export function getWeeklyTokenLimitWindowStart(now = new Date()) {
|
||||
const current = toValidDate(now);
|
||||
if (!current) throw new Error("A valid current time is required");
|
||||
|
||||
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`);
|
||||
}
|
||||
|
||||
export function getSessionResetAt(sessionStartedAt) {
|
||||
const sessionStart = toValidDate(sessionStartedAt);
|
||||
return sessionStart
|
||||
? new Date(sessionStart.getTime() + USER_TOKEN_LIMIT_SESSION_MS)
|
||||
: null;
|
||||
}
|
||||
|
||||
export function getActiveSessionWindowStart(sessionStartedAt, now = new Date()) {
|
||||
const current = toValidDate(now);
|
||||
const sessionResetAt = getSessionResetAt(sessionStartedAt);
|
||||
if (!current || !sessionResetAt || sessionResetAt.getTime() <= current.getTime()) return null;
|
||||
return toValidDate(sessionStartedAt);
|
||||
}
|
||||
+59
-54
@@ -1,13 +1,18 @@
|
||||
import {
|
||||
ensureUserTokenQuotaSession,
|
||||
getUserProviderEarliestTokenUsageSince,
|
||||
getUserProviderTokenUsageSince,
|
||||
getUserTokenQuotaSession,
|
||||
getUserTokenLimits,
|
||||
} from "@/lib/db/index.js";
|
||||
import { getUserTokenLimitWindowStart } from "@/lib/tokenLimitEnforcer.js";
|
||||
import {
|
||||
getActiveSessionWindowStart,
|
||||
getRollingSessionWindowStart,
|
||||
getSessionResetAt,
|
||||
getWeeklyTokenLimitWindowStart,
|
||||
} from "@/lib/userTokenLimitWindows.js";
|
||||
import {
|
||||
USER_TOKEN_LIMIT_PROVIDER_IDS,
|
||||
USER_TOKEN_LIMIT_WINDOW_IDS,
|
||||
USER_TOKEN_LIMIT_SESSION_MS,
|
||||
USER_TOKEN_LIMIT_WEEKLY_MS,
|
||||
USER_TOKEN_LIMIT_WINDOWS,
|
||||
} from "open-sse/config/userTokenLimits.js";
|
||||
@@ -37,19 +42,33 @@ export function buildUserTokenQuotaWindow(limit, used, windowStart) {
|
||||
};
|
||||
}
|
||||
|
||||
function getSessionNextTokenRestoreAt(earliestTokenUsageAt, now) {
|
||||
if (!earliestTokenUsageAt) return null;
|
||||
|
||||
const expiryTime = new Date(earliestTokenUsageAt).getTime() + USER_TOKEN_LIMIT_SESSION_MS;
|
||||
return Number.isFinite(expiryTime) && expiryTime > now.getTime()
|
||||
? new Date(expiryTime).toISOString()
|
||||
: null;
|
||||
}
|
||||
|
||||
function getWeeklyResetAt(windowStart) {
|
||||
return new Date(windowStart.getTime() + USER_TOKEN_LIMIT_WEEKLY_MS).toISOString();
|
||||
}
|
||||
|
||||
async function getActiveSessionStart(userId, provider, now) {
|
||||
const storedSessionStart = await getUserTokenQuotaSession(userId, provider);
|
||||
if (storedSessionStart) {
|
||||
return getActiveSessionWindowStart(storedSessionStart, now);
|
||||
}
|
||||
|
||||
// Seed a fixed-session record for usage logged before the session table.
|
||||
const earliestUsageAt = await getUserProviderEarliestTokenUsageSince(
|
||||
userId,
|
||||
provider,
|
||||
getRollingSessionWindowStart(now),
|
||||
);
|
||||
const legacySessionStart = getActiveSessionWindowStart(earliestUsageAt, now);
|
||||
if (!legacySessionStart) return null;
|
||||
|
||||
const savedSessionStart = await ensureUserTokenQuotaSession(
|
||||
userId,
|
||||
provider,
|
||||
legacySessionStart,
|
||||
);
|
||||
return getActiveSessionWindowStart(savedSessionStart, now);
|
||||
}
|
||||
|
||||
/**
|
||||
* Return the configured token-budget usage for a dashboard user.
|
||||
* Limits of zero intentionally remain unlimited while still reporting use.
|
||||
@@ -59,53 +78,39 @@ export async function getUserTokenQuota(userId, now = new Date()) {
|
||||
throw new Error("User id is required");
|
||||
}
|
||||
|
||||
const windows = Object.fromEntries(USER_TOKEN_LIMIT_WINDOW_IDS.map((windowType) => [
|
||||
windowType,
|
||||
getUserTokenLimitWindowStart(windowType, now),
|
||||
]));
|
||||
const limits = await getUserTokenLimits(userId);
|
||||
const weeklyWindowStart = getWeeklyTokenLimitWindowStart(now);
|
||||
const weeklyResetAt = getWeeklyResetAt(weeklyWindowStart);
|
||||
|
||||
const usageEntries = await Promise.all(
|
||||
USER_TOKEN_LIMIT_PROVIDER_IDS.flatMap((provider) => (
|
||||
USER_TOKEN_LIMIT_WINDOW_IDS.map(async (windowType) => [
|
||||
provider,
|
||||
windowType,
|
||||
await getUserProviderTokenUsageSince(userId, provider, windows[windowType]),
|
||||
])
|
||||
)),
|
||||
);
|
||||
const providerEntries = await Promise.all(USER_TOKEN_LIMIT_PROVIDER_IDS.map(async (provider) => {
|
||||
const sessionWindowStart = await getActiveSessionStart(userId, provider, now);
|
||||
const [sessionUsed, weeklyUsed] = await Promise.all([
|
||||
sessionWindowStart
|
||||
? getUserProviderTokenUsageSince(userId, provider, sessionWindowStart)
|
||||
: 0,
|
||||
getUserProviderTokenUsageSince(userId, provider, weeklyWindowStart),
|
||||
]);
|
||||
|
||||
const sessionTokenUsageEntries = await Promise.all(
|
||||
USER_TOKEN_LIMIT_PROVIDER_IDS.map(async (provider) => [
|
||||
provider,
|
||||
await getUserProviderEarliestTokenUsageSince(
|
||||
userId,
|
||||
provider,
|
||||
windows[USER_TOKEN_LIMIT_WINDOWS.SESSION],
|
||||
),
|
||||
]),
|
||||
const session = buildUserTokenQuotaWindow(
|
||||
limits[provider]?.[USER_TOKEN_LIMIT_WINDOWS.SESSION],
|
||||
sessionUsed,
|
||||
sessionWindowStart || now,
|
||||
);
|
||||
const sessionNextTokenRestoreAt = Object.fromEntries(
|
||||
sessionTokenUsageEntries.map(([provider, timestamp]) => [
|
||||
provider,
|
||||
getSessionNextTokenRestoreAt(timestamp, now),
|
||||
]),
|
||||
);
|
||||
const weeklyResetAt = getWeeklyResetAt(windows[USER_TOKEN_LIMIT_WINDOWS.WEEKLY]);
|
||||
session.windowStart = sessionWindowStart?.toISOString() || null;
|
||||
session.resetAt = getSessionResetAt(sessionWindowStart)?.toISOString() || null;
|
||||
|
||||
const providers = Object.fromEntries(
|
||||
USER_TOKEN_LIMIT_PROVIDER_IDS.map((provider) => [provider, {}]),
|
||||
const weekly = buildUserTokenQuotaWindow(
|
||||
limits[provider]?.[USER_TOKEN_LIMIT_WINDOWS.WEEKLY],
|
||||
weeklyUsed,
|
||||
weeklyWindowStart,
|
||||
);
|
||||
for (const [provider, windowType, used] of usageEntries) {
|
||||
providers[provider][windowType] = buildUserTokenQuotaWindow(
|
||||
limits[provider]?.[windowType],
|
||||
used,
|
||||
windows[windowType],
|
||||
);
|
||||
providers[provider][windowType].resetAt = windowType === USER_TOKEN_LIMIT_WINDOWS.SESSION
|
||||
? sessionNextTokenRestoreAt[provider]
|
||||
: weeklyResetAt;
|
||||
}
|
||||
weekly.resetAt = weeklyResetAt;
|
||||
|
||||
return providers;
|
||||
return [provider, {
|
||||
[USER_TOKEN_LIMIT_WINDOWS.SESSION]: session,
|
||||
[USER_TOKEN_LIMIT_WINDOWS.WEEKLY]: weekly,
|
||||
}];
|
||||
}));
|
||||
|
||||
return Object.fromEntries(providerEntries);
|
||||
}
|
||||
|
||||
@@ -2,11 +2,17 @@ import { beforeEach, describe, expect, it, vi } from "vitest";
|
||||
|
||||
const getUserById = vi.fn();
|
||||
const getUserProviderTokenUsageSince = vi.fn();
|
||||
const getUserProviderEarliestTokenUsageSince = vi.fn();
|
||||
const getUserTokenQuotaSession = vi.fn();
|
||||
const ensureUserTokenQuotaSession = vi.fn();
|
||||
const getUserTokenLimits = vi.fn();
|
||||
|
||||
vi.mock("@/lib/db/index.js", () => ({
|
||||
getUserById,
|
||||
ensureUserTokenQuotaSession,
|
||||
getUserProviderEarliestTokenUsageSince,
|
||||
getUserProviderTokenUsageSince,
|
||||
getUserTokenQuotaSession,
|
||||
getUserTokenLimits,
|
||||
}));
|
||||
|
||||
@@ -18,13 +24,17 @@ const {
|
||||
describe("user token limit enforcement", () => {
|
||||
beforeEach(() => {
|
||||
getUserById.mockReset();
|
||||
ensureUserTokenQuotaSession.mockReset();
|
||||
getUserProviderEarliestTokenUsageSince.mockReset();
|
||||
getUserProviderTokenUsageSince.mockReset();
|
||||
getUserTokenQuotaSession.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 },
|
||||
});
|
||||
getUserTokenQuotaSession.mockResolvedValue("2026-07-17T06:00:00.000Z");
|
||||
});
|
||||
|
||||
it("calculates rolling session and Monday Vietnam weekly window starts", () => {
|
||||
@@ -55,6 +65,23 @@ describe("user token limit enforcement", () => {
|
||||
expect(getUserProviderTokenUsageSince).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
it("does not block from an expired fixed session until a new request begins one", async () => {
|
||||
getUserTokenQuotaSession.mockResolvedValue("2026-07-17T05:00:00.000Z");
|
||||
getUserProviderTokenUsageSince.mockResolvedValueOnce(0);
|
||||
|
||||
await expect(checkUserTokenLimit(
|
||||
"user-1",
|
||||
"orbit-provider",
|
||||
new Date("2026-07-17T10:10:00.000Z"),
|
||||
)).resolves.toBeNull();
|
||||
|
||||
expect(getUserProviderTokenUsageSince).toHaveBeenCalledWith(
|
||||
"user-1",
|
||||
"orbit-provider",
|
||||
new Date("2026-07-12T17:00:00.000Z"),
|
||||
);
|
||||
});
|
||||
|
||||
it("checks weekly usage after the session window still has headroom", async () => {
|
||||
getUserProviderTokenUsageSince
|
||||
.mockResolvedValueOnce(80)
|
||||
|
||||
@@ -67,6 +67,31 @@ describe("user token limit repository", () => {
|
||||
)).resolves.toBe(50);
|
||||
});
|
||||
|
||||
it("keeps a fixed session anchor until its five-hour boundary passes", async () => {
|
||||
const db = await import("@/lib/db/index.js");
|
||||
const user = await db.createUser({ username: "session-user", password: "password", role: "user" });
|
||||
|
||||
await expect(db.ensureUserTokenQuotaSession(
|
||||
user.id,
|
||||
"codex",
|
||||
"2026-07-17T06:00:00.000Z",
|
||||
)).resolves.toBe("2026-07-17T06:00:00.000Z");
|
||||
|
||||
await expect(db.ensureUserTokenQuotaSession(
|
||||
user.id,
|
||||
"codex",
|
||||
"2026-07-17T08:00:00.000Z",
|
||||
)).resolves.toBe("2026-07-17T06:00:00.000Z");
|
||||
|
||||
await expect(db.ensureUserTokenQuotaSession(
|
||||
user.id,
|
||||
"codex",
|
||||
"2026-07-17T11:01:00.000Z",
|
||||
)).resolves.toBe("2026-07-17T11:01:00.000Z");
|
||||
await expect(db.getUserTokenQuotaSession(user.id, "codex"))
|
||||
.resolves.toBe("2026-07-17T11:01:00.000Z");
|
||||
});
|
||||
|
||||
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" });
|
||||
|
||||
@@ -3,19 +3,21 @@ import { beforeEach, describe, expect, it, vi } from "vitest";
|
||||
const getUserTokenLimits = vi.fn();
|
||||
const getUserProviderTokenUsageSince = vi.fn();
|
||||
const getUserProviderEarliestTokenUsageSince = vi.fn();
|
||||
const getUserTokenLimitWindowStart = vi.fn();
|
||||
const getUserTokenQuotaSession = vi.fn();
|
||||
const ensureUserTokenQuotaSession = vi.fn();
|
||||
|
||||
vi.mock("@/lib/db/index.js", () => ({
|
||||
getUserTokenLimits,
|
||||
getUserProviderTokenUsageSince,
|
||||
getUserProviderEarliestTokenUsageSince,
|
||||
getUserTokenQuotaSession,
|
||||
ensureUserTokenQuotaSession,
|
||||
}));
|
||||
vi.mock("@/lib/tokenLimitEnforcer.js", () => ({ getUserTokenLimitWindowStart }));
|
||||
|
||||
const { getUserTokenQuota } = await import("@/lib/userTokenQuota.js");
|
||||
|
||||
const sessionStart = new Date("2026-07-17T05:00:00.000Z");
|
||||
const weeklyStart = new Date("2026-07-13T17:00:00.000Z");
|
||||
const sessionStart = new Date("2026-07-17T06:00:00.000Z");
|
||||
const weeklyStart = new Date("2026-07-12T17:00:00.000Z");
|
||||
|
||||
function usageKey(provider, windowType) {
|
||||
return `${provider}:${windowType}`;
|
||||
@@ -26,15 +28,14 @@ describe("user token quota snapshot", () => {
|
||||
getUserTokenLimits.mockReset();
|
||||
getUserProviderTokenUsageSince.mockReset();
|
||||
getUserProviderEarliestTokenUsageSince.mockReset();
|
||||
getUserTokenLimitWindowStart.mockReset();
|
||||
getUserTokenQuotaSession.mockReset();
|
||||
ensureUserTokenQuotaSession.mockReset();
|
||||
|
||||
getUserTokenLimits.mockResolvedValue({
|
||||
"orbit-provider": { session: 100, weekly: 1000 },
|
||||
codex: { session: 0, weekly: 500 },
|
||||
});
|
||||
getUserTokenLimitWindowStart.mockImplementation((windowType) => (
|
||||
windowType === "session" ? sessionStart : weeklyStart
|
||||
));
|
||||
getUserTokenQuotaSession.mockResolvedValue(sessionStart.toISOString());
|
||||
const usage = new Map([
|
||||
[usageKey("orbit-provider", "session"), 25],
|
||||
[usageKey("orbit-provider", "weekly"), 1200],
|
||||
@@ -62,21 +63,44 @@ describe("user token quota snapshot", () => {
|
||||
},
|
||||
});
|
||||
expect(quota.codex.session.windowStart).toBe(sessionStart.toISOString());
|
||||
expect(quota.codex.session.resetAt).toBeNull();
|
||||
expect(quota.codex.weekly.resetAt).toBe("2026-07-20T17:00:00.000Z");
|
||||
expect(quota.codex.session.resetAt).toBe("2026-07-17T11:00:00.000Z");
|
||||
expect(quota.codex.weekly.resetAt).toBe("2026-07-19T17:00:00.000Z");
|
||||
expect(getUserProviderTokenUsageSince).toHaveBeenCalledTimes(4);
|
||||
expect(getUserProviderEarliestTokenUsageSince).toHaveBeenCalledTimes(2);
|
||||
expect(getUserProviderEarliestTokenUsageSince).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("reports when the next tokens leave a rolling session window", async () => {
|
||||
it("resets an expired session to its full budget until another request starts one", async () => {
|
||||
getUserTokenQuotaSession.mockResolvedValue("2026-07-17T05:00:00.000Z");
|
||||
|
||||
const quota = await getUserTokenQuota("user-1", new Date("2026-07-17T10:10:00.000Z"));
|
||||
|
||||
expect(quota["orbit-provider"].session).toMatchObject({
|
||||
used: 0,
|
||||
remaining: 100,
|
||||
remainingPercentage: 100,
|
||||
windowStart: null,
|
||||
resetAt: null,
|
||||
});
|
||||
expect(quota.codex.session).toMatchObject({ used: 0, remaining: null, windowStart: null, resetAt: null });
|
||||
expect(getUserProviderEarliestTokenUsageSince).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("seeds an active fixed session from legacy usage exactly once", async () => {
|
||||
getUserTokenQuotaSession.mockResolvedValue(null);
|
||||
getUserProviderEarliestTokenUsageSince.mockImplementation(async (_userId, provider) => (
|
||||
provider === "orbit-provider" ? "2026-07-17T06:30:00.000Z" : null
|
||||
));
|
||||
ensureUserTokenQuotaSession.mockResolvedValue("2026-07-17T06:30:00.000Z");
|
||||
|
||||
const quota = await getUserTokenQuota("user-1", new Date("2026-07-17T10:00:00.000Z"));
|
||||
|
||||
expect(quota["orbit-provider"].session.resetAt).toBe("2026-07-17T11:30:00.000Z");
|
||||
expect(quota.codex.session.resetAt).toBeNull();
|
||||
expect(ensureUserTokenQuotaSession).toHaveBeenCalledWith(
|
||||
"user-1",
|
||||
"orbit-provider",
|
||||
new Date("2026-07-17T06:30:00.000Z"),
|
||||
);
|
||||
});
|
||||
|
||||
it("requires a user id", async () => {
|
||||
|
||||
Reference in New Issue
Block a user