fix: logic reset quota by session

This commit is contained in:
2026-07-20 16:18:00 +07:00
parent 0e8a21fdbe
commit 580bf9d9ae
16 changed files with 340 additions and 87 deletions
+3
View File
@@ -1,5 +1,8 @@
# v0.5.35 (2026-07-16) # 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 ## Features
- **User limits**: add per-user total-token budgets for Orbit Provider and Codex with rolling 5-hour and weekly windows - **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 - **User quota**: show remaining Orbit and Codex headroom in the users table with session and weekly usage details
+1 -1
View File
@@ -19,7 +19,7 @@ export const USER_TOKEN_LIMIT_WINDOW_IDS = Object.freeze(
export const USER_TOKEN_LIMIT_WINDOW_CONFIG = Object.freeze({ export const USER_TOKEN_LIMIT_WINDOW_CONFIG = Object.freeze({
[USER_TOKEN_LIMIT_WINDOWS.SESSION]: Object.freeze({ [USER_TOKEN_LIMIT_WINDOWS.SESSION]: Object.freeze({
name: "Session", name: "Session",
description: "Rolling 5 hours", description: "Fixed 5 hours from the first request",
}), }),
[USER_TOKEN_LIMIT_WINDOWS.WEEKLY]: Object.freeze({ [USER_TOKEN_LIMIT_WINDOWS.WEEKLY]: Object.freeze({
name: "Weekly", name: "Weekly",
@@ -36,8 +36,6 @@ function getQuotaTone(percentage) {
return { bar: "bg-red-500", dot: "bg-red-500", text: "text-red-500" }; return { bar: "bg-red-500", dot: "bg-red-500", text: "text-red-500" };
} }
const INACTIVE_SESSION_COUNTDOWN = "4h59m";
function TokenQuotaResetStatus({ quota }) { function TokenQuotaResetStatus({ quota }) {
const [now, setNow] = useState(() => new Date()); const [now, setNow] = useState(() => new Date());
@@ -51,7 +49,7 @@ function TokenQuotaResetStatus({ quota }) {
const countdown = formatResetTime(quota.resetAt, now); const countdown = formatResetTime(quota.resetAt, now);
const isSession = quota.windowType === USER_TOKEN_LIMIT_WINDOWS.SESSION; const isSession = quota.windowType === USER_TOKEN_LIMIT_WINDOWS.SESSION;
const text = countdown === "-" 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}`); : (isSession ? `Next tokens restore in ${countdown}` : `Resets in ${countdown}`);
return ( return (
+1 -1
View File
@@ -358,7 +358,7 @@ export default function UsersPage() {
<span className="material-symbols-outlined mt-0.5 text-[20px] text-brand-500">hourglass_top</span> <span className="material-symbols-outlined mt-0.5 text-[20px] text-brand-500">hourglass_top</span>
<div> <div>
<p className="text-sm font-medium text-text-main">Total token budgets</p> <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> </div>
</div> </div>
+2 -1
View File
@@ -63,7 +63,8 @@ export {
export { export {
createEmptyUserTokenLimits, getUserTokenLimits, createEmptyUserTokenLimits, getUserTokenLimits,
replaceUserTokenLimits, getUserProviderTokenUsageSince, replaceUserTokenLimits, getUserProviderTokenUsageSince,
getUserProviderEarliestTokenUsageSince, getUserProviderEarliestTokenUsageSince, getUserTokenQuotaSession,
ensureUserTokenQuotaSession,
} from "./repos/userTokenLimitsRepo.js"; } from "./repos/userTokenLimitsRepo.js";
// Aliases (model + custom + mitm) // 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;
+2 -1
View File
@@ -10,8 +10,9 @@ import m006 from "./006-combo-owners.js";
import m007 from "./007-admin-provider-connections.js"; import m007 from "./007-admin-provider-connections.js";
import m008 from "./008-user-token-limits.js"; import m008 from "./008-user-token-limits.js";
import m009 from "./009-remove-tunnel-settings.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() { export function latestVersion() {
return MIGRATIONS.length ? MIGRATIONS[MIGRATIONS.length - 1].version : 0; return MIGRATIONS.length ? MIGRATIONS[MIGRATIONS.length - 1].version : 0;
+9
View File
@@ -3,6 +3,8 @@ import { getAdapter } from "../driver.js";
import { parseJson, stringifyJson } from "../helpers/jsonCol.js"; import { parseJson, stringifyJson } from "../helpers/jsonCol.js";
import { getMeta, setMeta } from "../helpers/metaStore.js"; import { getMeta, setMeta } from "../helpers/metaStore.js";
import { appendUsageAccessClause, getUsageAccessScope } from "./usageAccessScope.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 { import {
formatVietnamDateTime, formatVietnamDateTime,
getVietnamDateKey, getVietnamDateKey,
@@ -20,6 +22,7 @@ const PENDING_TIMEOUT_MS = 60 * 1000;
const RING_CAP = 50; const RING_CAP = 50;
const CONN_CACHE_TTL_MS = 30 * 1000; const CONN_CACHE_TTL_MS = 30 * 1000;
const PERIOD_MS = { "24h": 86400000, "7d": 604800000, "30d": 2592000000, "60d": 5184000000 }; 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 // In-memory state shared across Next.js modules
if (!global._pendingRequests) global._pendingRequests = { byModel: {}, byAccount: {} }; if (!global._pendingRequests) global._pendingRequests = { byModel: {}, byAccount: {} };
@@ -316,6 +319,7 @@ export async function saveRequestUsage(entry) {
stringifyJson(tokens), stringifyJson({}), stringifyJson(tokens), stringifyJson({}),
] ]
); );
inserted = true;
const dateKey = getLocalDateKey(entry.timestamp); const dateKey = getLocalDateKey(entry.timestamp);
const row = db.get(`SELECT data FROM usageDaily WHERE dateKey = ?`, [dateKey]); 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, requests: 0, promptTokens: 0, completionTokens: 0, cost: 0,
byProvider: {}, byModel: {}, byAccount: {}, byApiKey: {}, byUser: {}, byEndpoint: {}, byProvider: {}, byModel: {}, byAccount: {}, byApiKey: {}, byUser: {}, byEndpoint: {},
}; };
aggregateEntryToDay(day, entry); aggregateEntryToDay(day, entry);
db.run(`INSERT INTO usageDaily(dateKey, data) VALUES(?, ?) ON CONFLICT(dateKey) DO UPDATE SET data = excluded.data`, [dateKey, stringifyJson(day)]); 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; inserted = true;
}); });
if (inserted && entry.userId && userTokenQuotaProviderSet.has(entry.provider) && (promptTokens + completionTokens) > 0) {
await ensureUserTokenQuotaSession(entry.userId, entry.provider, entry.timestamp);
}
if (inserted) { if (inserted) {
pushToRing(entry); pushToRing(entry);
scheduleStatsEvent("update", 250); scheduleStatsEvent("update", 250);
+65
View File
@@ -1,6 +1,7 @@
import { getAdapter } from "../driver.js"; import { getAdapter } from "../driver.js";
import { import {
USER_TOKEN_LIMIT_PROVIDER_IDS, USER_TOKEN_LIMIT_PROVIDER_IDS,
USER_TOKEN_LIMIT_SESSION_MS,
USER_TOKEN_LIMIT_WINDOW_IDS, USER_TOKEN_LIMIT_WINDOW_IDS,
} from "open-sse/config/userTokenLimits.js"; } from "open-sse/config/userTokenLimits.js";
@@ -136,4 +137,68 @@ export async function getUserProviderEarliestTokenUsageSince(userId, provider, s
[userId, provider, since.toISOString()], [userId, provider, since.toISOString()],
); );
return row?.timestamp || null; 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;
} }
+12
View File
@@ -148,6 +148,18 @@ export const TABLES = {
"CREATE INDEX IF NOT EXISTS idx_user_token_limits_user ON userTokenLimits(userId)", "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: { kv: {
columns: { columns: {
scope: "TEXT NOT NULL", scope: "TEXT NOT NULL",
+38 -12
View File
@@ -1,17 +1,20 @@
import { import {
USER_TOKEN_LIMIT_PROVIDER_IDS, USER_TOKEN_LIMIT_PROVIDER_IDS,
USER_TOKEN_LIMIT_SESSION_MS,
USER_TOKEN_LIMIT_WINDOWS, USER_TOKEN_LIMIT_WINDOWS,
} from "open-sse/config/userTokenLimits.js"; } from "open-sse/config/userTokenLimits.js";
import { import {
ensureUserTokenQuotaSession,
getUserById, getUserById,
getUserProviderEarliestTokenUsageSince,
getUserProviderTokenUsageSince, getUserProviderTokenUsageSince,
getUserTokenQuotaSession,
getUserTokenLimits, getUserTokenLimits,
} from "@/lib/db/index.js"; } from "@/lib/db/index.js";
import { import {
getVietnamDateKey, getActiveSessionWindowStart,
shiftVietnamDateKey, getRollingSessionWindowStart,
} from "@/shared/utils/dateTime.js"; getWeeklyTokenLimitWindowStart,
} from "@/lib/userTokenLimitWindows.js";
const limitedProviderSet = new Set(USER_TOKEN_LIMIT_PROVIDER_IDS); 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 (!Number.isFinite(current.getTime())) throw new Error("A valid current time is required");
if (windowType === USER_TOKEN_LIMIT_WINDOWS.SESSION) { 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) { if (windowType === USER_TOKEN_LIMIT_WINDOWS.WEEKLY) {
const dateKey = getVietnamDateKey(current); return getWeeklyTokenLimitWindowStart(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"); 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. * 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. * 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]; const limit = providerLimits[windowType];
if (!Number.isSafeInteger(limit) || limit <= 0) continue; if (!Number.isSafeInteger(limit) || limit <= 0) continue;
const windowStart = getUserTokenLimitWindowStart(windowType, now); const windowStart = windowType === USER_TOKEN_LIMIT_WINDOWS.SESSION
const used = await getUserProviderTokenUsageSince(user.id, provider, windowStart); ? await getActiveSessionStart(user.id, provider, now)
: getWeeklyTokenLimitWindowStart(now);
const used = windowStart
? await getUserProviderTokenUsageSince(user.id, provider, windowStart)
: 0;
if (used >= limit) { if (used >= limit) {
return { return {
exceeded: true, exceeded: true,
+44
View File
@@ -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);
}
+61 -56
View File
@@ -1,13 +1,18 @@
import { import {
ensureUserTokenQuotaSession,
getUserProviderEarliestTokenUsageSince, getUserProviderEarliestTokenUsageSince,
getUserProviderTokenUsageSince, getUserProviderTokenUsageSince,
getUserTokenQuotaSession,
getUserTokenLimits, getUserTokenLimits,
} from "@/lib/db/index.js"; } from "@/lib/db/index.js";
import { getUserTokenLimitWindowStart } from "@/lib/tokenLimitEnforcer.js"; import {
getActiveSessionWindowStart,
getRollingSessionWindowStart,
getSessionResetAt,
getWeeklyTokenLimitWindowStart,
} from "@/lib/userTokenLimitWindows.js";
import { import {
USER_TOKEN_LIMIT_PROVIDER_IDS, USER_TOKEN_LIMIT_PROVIDER_IDS,
USER_TOKEN_LIMIT_WINDOW_IDS,
USER_TOKEN_LIMIT_SESSION_MS,
USER_TOKEN_LIMIT_WEEKLY_MS, USER_TOKEN_LIMIT_WEEKLY_MS,
USER_TOKEN_LIMIT_WINDOWS, USER_TOKEN_LIMIT_WINDOWS,
} from "open-sse/config/userTokenLimits.js"; } 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) { function getWeeklyResetAt(windowStart) {
return new Date(windowStart.getTime() + USER_TOKEN_LIMIT_WEEKLY_MS).toISOString(); 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. * Return the configured token-budget usage for a dashboard user.
* Limits of zero intentionally remain unlimited while still reporting use. * 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"); 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 limits = await getUserTokenLimits(userId);
const weeklyWindowStart = getWeeklyTokenLimitWindowStart(now);
const weeklyResetAt = getWeeklyResetAt(weeklyWindowStart);
const usageEntries = await Promise.all( const providerEntries = await Promise.all(USER_TOKEN_LIMIT_PROVIDER_IDS.map(async (provider) => {
USER_TOKEN_LIMIT_PROVIDER_IDS.flatMap((provider) => ( const sessionWindowStart = await getActiveSessionStart(userId, provider, now);
USER_TOKEN_LIMIT_WINDOW_IDS.map(async (windowType) => [ const [sessionUsed, weeklyUsed] = await Promise.all([
provider, sessionWindowStart
windowType, ? getUserProviderTokenUsageSince(userId, provider, sessionWindowStart)
await getUserProviderTokenUsageSince(userId, provider, windows[windowType]), : 0,
]) getUserProviderTokenUsageSince(userId, provider, weeklyWindowStart),
)), ]);
);
const sessionTokenUsageEntries = await Promise.all( const session = buildUserTokenQuotaWindow(
USER_TOKEN_LIMIT_PROVIDER_IDS.map(async (provider) => [ limits[provider]?.[USER_TOKEN_LIMIT_WINDOWS.SESSION],
provider, sessionUsed,
await getUserProviderEarliestTokenUsageSince( sessionWindowStart || now,
userId,
provider,
windows[USER_TOKEN_LIMIT_WINDOWS.SESSION],
),
]),
);
const sessionNextTokenRestoreAt = Object.fromEntries(
sessionTokenUsageEntries.map(([provider, timestamp]) => [
provider,
getSessionNextTokenRestoreAt(timestamp, now),
]),
);
const weeklyResetAt = getWeeklyResetAt(windows[USER_TOKEN_LIMIT_WINDOWS.WEEKLY]);
const providers = Object.fromEntries(
USER_TOKEN_LIMIT_PROVIDER_IDS.map((provider) => [provider, {}]),
);
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 session.windowStart = sessionWindowStart?.toISOString() || null;
? sessionNextTokenRestoreAt[provider] session.resetAt = getSessionResetAt(sessionWindowStart)?.toISOString() || null;
: weeklyResetAt;
}
return providers; const weekly = buildUserTokenQuotaWindow(
limits[provider]?.[USER_TOKEN_LIMIT_WINDOWS.WEEKLY],
weeklyUsed,
weeklyWindowStart,
);
weekly.resetAt = weeklyResetAt;
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 getUserById = vi.fn();
const getUserProviderTokenUsageSince = vi.fn(); const getUserProviderTokenUsageSince = vi.fn();
const getUserProviderEarliestTokenUsageSince = vi.fn();
const getUserTokenQuotaSession = vi.fn();
const ensureUserTokenQuotaSession = vi.fn();
const getUserTokenLimits = vi.fn(); const getUserTokenLimits = vi.fn();
vi.mock("@/lib/db/index.js", () => ({ vi.mock("@/lib/db/index.js", () => ({
getUserById, getUserById,
ensureUserTokenQuotaSession,
getUserProviderEarliestTokenUsageSince,
getUserProviderTokenUsageSince, getUserProviderTokenUsageSince,
getUserTokenQuotaSession,
getUserTokenLimits, getUserTokenLimits,
})); }));
@@ -18,13 +24,17 @@ const {
describe("user token limit enforcement", () => { describe("user token limit enforcement", () => {
beforeEach(() => { beforeEach(() => {
getUserById.mockReset(); getUserById.mockReset();
ensureUserTokenQuotaSession.mockReset();
getUserProviderEarliestTokenUsageSince.mockReset();
getUserProviderTokenUsageSince.mockReset(); getUserProviderTokenUsageSince.mockReset();
getUserTokenQuotaSession.mockReset();
getUserTokenLimits.mockReset(); getUserTokenLimits.mockReset();
getUserById.mockResolvedValue({ id: "user-1", role: "user", isActive: true }); getUserById.mockResolvedValue({ id: "user-1", role: "user", isActive: true });
getUserTokenLimits.mockResolvedValue({ getUserTokenLimits.mockResolvedValue({
"orbit-provider": { session: 100, weekly: 1000 }, "orbit-provider": { session: 100, weekly: 1000 },
codex: { session: 200, weekly: 2000 }, codex: { session: 200, weekly: 2000 },
}); });
getUserTokenQuotaSession.mockResolvedValue("2026-07-17T06:00:00.000Z");
}); });
it("calculates rolling session and Monday Vietnam weekly window starts", () => { it("calculates rolling session and Monday Vietnam weekly window starts", () => {
@@ -55,6 +65,23 @@ describe("user token limit enforcement", () => {
expect(getUserProviderTokenUsageSince).toHaveBeenCalledTimes(1); 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 () => { it("checks weekly usage after the session window still has headroom", async () => {
getUserProviderTokenUsageSince getUserProviderTokenUsageSince
.mockResolvedValueOnce(80) .mockResolvedValueOnce(80)
+25
View File
@@ -67,6 +67,31 @@ describe("user token limit repository", () => {
)).resolves.toBe(50); )).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 () => { it("rejects negative and non-integer limits without changing stored values", async () => {
const db = await import("@/lib/db/index.js"); const db = await import("@/lib/db/index.js");
const user = await db.createUser({ username: "invalid-limit", password: "password", role: "user" }); const user = await db.createUser({ username: "invalid-limit", password: "password", role: "user" });
+36 -12
View File
@@ -3,19 +3,21 @@ import { beforeEach, describe, expect, it, vi } from "vitest";
const getUserTokenLimits = vi.fn(); const getUserTokenLimits = vi.fn();
const getUserProviderTokenUsageSince = vi.fn(); const getUserProviderTokenUsageSince = vi.fn();
const getUserProviderEarliestTokenUsageSince = vi.fn(); const getUserProviderEarliestTokenUsageSince = vi.fn();
const getUserTokenLimitWindowStart = vi.fn(); const getUserTokenQuotaSession = vi.fn();
const ensureUserTokenQuotaSession = vi.fn();
vi.mock("@/lib/db/index.js", () => ({ vi.mock("@/lib/db/index.js", () => ({
getUserTokenLimits, getUserTokenLimits,
getUserProviderTokenUsageSince, getUserProviderTokenUsageSince,
getUserProviderEarliestTokenUsageSince, getUserProviderEarliestTokenUsageSince,
getUserTokenQuotaSession,
ensureUserTokenQuotaSession,
})); }));
vi.mock("@/lib/tokenLimitEnforcer.js", () => ({ getUserTokenLimitWindowStart }));
const { getUserTokenQuota } = await import("@/lib/userTokenQuota.js"); const { getUserTokenQuota } = await import("@/lib/userTokenQuota.js");
const sessionStart = new Date("2026-07-17T05:00:00.000Z"); const sessionStart = new Date("2026-07-17T06:00:00.000Z");
const weeklyStart = new Date("2026-07-13T17:00:00.000Z"); const weeklyStart = new Date("2026-07-12T17:00:00.000Z");
function usageKey(provider, windowType) { function usageKey(provider, windowType) {
return `${provider}:${windowType}`; return `${provider}:${windowType}`;
@@ -26,15 +28,14 @@ describe("user token quota snapshot", () => {
getUserTokenLimits.mockReset(); getUserTokenLimits.mockReset();
getUserProviderTokenUsageSince.mockReset(); getUserProviderTokenUsageSince.mockReset();
getUserProviderEarliestTokenUsageSince.mockReset(); getUserProviderEarliestTokenUsageSince.mockReset();
getUserTokenLimitWindowStart.mockReset(); getUserTokenQuotaSession.mockReset();
ensureUserTokenQuotaSession.mockReset();
getUserTokenLimits.mockResolvedValue({ getUserTokenLimits.mockResolvedValue({
"orbit-provider": { session: 100, weekly: 1000 }, "orbit-provider": { session: 100, weekly: 1000 },
codex: { session: 0, weekly: 500 }, codex: { session: 0, weekly: 500 },
}); });
getUserTokenLimitWindowStart.mockImplementation((windowType) => ( getUserTokenQuotaSession.mockResolvedValue(sessionStart.toISOString());
windowType === "session" ? sessionStart : weeklyStart
));
const usage = new Map([ const usage = new Map([
[usageKey("orbit-provider", "session"), 25], [usageKey("orbit-provider", "session"), 25],
[usageKey("orbit-provider", "weekly"), 1200], [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.windowStart).toBe(sessionStart.toISOString());
expect(quota.codex.session.resetAt).toBeNull(); expect(quota.codex.session.resetAt).toBe("2026-07-17T11:00:00.000Z");
expect(quota.codex.weekly.resetAt).toBe("2026-07-20T17:00:00.000Z"); expect(quota.codex.weekly.resetAt).toBe("2026-07-19T17:00:00.000Z");
expect(getUserProviderTokenUsageSince).toHaveBeenCalledTimes(4); 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) => ( getUserProviderEarliestTokenUsageSince.mockImplementation(async (_userId, provider) => (
provider === "orbit-provider" ? "2026-07-17T06:30:00.000Z" : null 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")); 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["orbit-provider"].session.resetAt).toBe("2026-07-17T11:30:00.000Z");
expect(quota.codex.session.resetAt).toBeNull(); 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 () => { it("requires a user id", async () => {