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)
## 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
+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({
[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 (
+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>
<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
View File
@@ -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;
+2 -1
View File
@@ -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;
+9
View File
@@ -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);
+65
View File
@@ -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;
}
+12
View File
@@ -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",
+38 -12
View File
@@ -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,
+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);
}
+59 -54
View File
@@ -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)
+25
View File
@@ -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" });
+36 -12
View File
@@ -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 () => {