From 580bf9d9ae42c137e57ee86c344592e195c76f58 Mon Sep 17 00:00:00 2001 From: Nezumi-2711 Date: Mon, 20 Jul 2026 16:18:00 +0700 Subject: [PATCH] fix: logic reset quota by session --- CHANGELOG.md | 3 + open-sse/config/userTokenLimits.js | 2 +- .../usage/components/SystemQuotaOverview.js | 4 +- src/app/(dashboard)/dashboard/users/page.js | 2 +- src/lib/db/index.js | 3 +- .../010-user-token-quota-sessions.js | 13 ++ src/lib/db/migrations/index.js | 3 +- src/lib/db/repos/usageRepo.js | 9 ++ src/lib/db/repos/userTokenLimitsRepo.js | 65 ++++++++++ src/lib/db/schema.js | 12 ++ src/lib/tokenLimitEnforcer.js | 50 ++++++-- src/lib/userTokenLimitWindows.js | 44 +++++++ src/lib/userTokenQuota.js | 117 +++++++++--------- tests/unit/user-token-limit-enforcer.test.js | 27 ++++ tests/unit/user-token-limits-repo.test.js | 25 ++++ tests/unit/user-token-quota.test.js | 48 +++++-- 16 files changed, 340 insertions(+), 87 deletions(-) create mode 100644 src/lib/db/migrations/010-user-token-quota-sessions.js create mode 100644 src/lib/userTokenLimitWindows.js diff --git a/CHANGELOG.md b/CHANGELOG.md index c56524cb..e080cb8e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -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 diff --git a/open-sse/config/userTokenLimits.js b/open-sse/config/userTokenLimits.js index 4cccd032..af4eecb1 100644 --- a/open-sse/config/userTokenLimits.js +++ b/open-sse/config/userTokenLimits.js @@ -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", diff --git a/src/app/(dashboard)/dashboard/usage/components/SystemQuotaOverview.js b/src/app/(dashboard)/dashboard/usage/components/SystemQuotaOverview.js index a1e6ebaa..620deddd 100644 --- a/src/app/(dashboard)/dashboard/usage/components/SystemQuotaOverview.js +++ b/src/app/(dashboard)/dashboard/usage/components/SystemQuotaOverview.js @@ -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 ( diff --git a/src/app/(dashboard)/dashboard/users/page.js b/src/app/(dashboard)/dashboard/users/page.js index 8eb492b6..1ae3f23f 100644 --- a/src/app/(dashboard)/dashboard/users/page.js +++ b/src/app/(dashboard)/dashboard/users/page.js @@ -358,7 +358,7 @@ export default function UsersPage() { 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.

+

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.

diff --git a/src/lib/db/index.js b/src/lib/db/index.js index 13965bc5..b6910a14 100644 --- a/src/lib/db/index.js +++ b/src/lib/db/index.js @@ -63,7 +63,8 @@ export { export { createEmptyUserTokenLimits, getUserTokenLimits, replaceUserTokenLimits, getUserProviderTokenUsageSince, - getUserProviderEarliestTokenUsageSince, + getUserProviderEarliestTokenUsageSince, getUserTokenQuotaSession, + ensureUserTokenQuotaSession, } from "./repos/userTokenLimitsRepo.js"; // Aliases (model + custom + mitm) diff --git a/src/lib/db/migrations/010-user-token-quota-sessions.js b/src/lib/db/migrations/010-user-token-quota-sessions.js new file mode 100644 index 00000000..8bc7f079 --- /dev/null +++ b/src/lib/db/migrations/010-user-token-quota-sessions.js @@ -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; \ No newline at end of file diff --git a/src/lib/db/migrations/index.js b/src/lib/db/migrations/index.js index 5571d7fb..abb1d21e 100644 --- a/src/lib/db/migrations/index.js +++ b/src/lib/db/migrations/index.js @@ -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; diff --git a/src/lib/db/repos/usageRepo.js b/src/lib/db/repos/usageRepo.js index 6acc0fcc..5336ca6f 100644 --- a/src/lib/db/repos/usageRepo.js +++ b/src/lib/db/repos/usageRepo.js @@ -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); diff --git a/src/lib/db/repos/userTokenLimitsRepo.js b/src/lib/db/repos/userTokenLimitsRepo.js index f146afe2..f9863610 100644 --- a/src/lib/db/repos/userTokenLimitsRepo.js +++ b/src/lib/db/repos/userTokenLimitsRepo.js @@ -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"; @@ -136,4 +137,68 @@ export async function getUserProviderEarliestTokenUsageSince(userId, provider, s [userId, provider, since.toISOString()], ); 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; } \ No newline at end of file diff --git a/src/lib/db/schema.js b/src/lib/db/schema.js index 87c73297..8e5fbd01 100644 --- a/src/lib/db/schema.js +++ b/src/lib/db/schema.js @@ -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", diff --git a/src/lib/tokenLimitEnforcer.js b/src/lib/tokenLimitEnforcer.js index 04fe6795..4ab68dc0 100644 --- a/src/lib/tokenLimitEnforcer.js +++ b/src/lib/tokenLimitEnforcer.js @@ -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, diff --git a/src/lib/userTokenLimitWindows.js b/src/lib/userTokenLimitWindows.js new file mode 100644 index 00000000..747558f3 --- /dev/null +++ b/src/lib/userTokenLimitWindows.js @@ -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); +} \ No newline at end of file diff --git a/src/lib/userTokenQuota.js b/src/lib/userTokenQuota.js index 531fae15..8e96e09e 100644 --- a/src/lib/userTokenQuota.js +++ b/src/lib/userTokenQuota.js @@ -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 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], + const session = buildUserTokenQuotaWindow( + limits[provider]?.[USER_TOKEN_LIMIT_WINDOWS.SESSION], + sessionUsed, + sessionWindowStart || now, ); - providers[provider][windowType].resetAt = windowType === USER_TOKEN_LIMIT_WINDOWS.SESSION - ? sessionNextTokenRestoreAt[provider] - : weeklyResetAt; - } + session.windowStart = sessionWindowStart?.toISOString() || null; + session.resetAt = getSessionResetAt(sessionWindowStart)?.toISOString() || null; - 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); } diff --git a/tests/unit/user-token-limit-enforcer.test.js b/tests/unit/user-token-limit-enforcer.test.js index b820fa73..66d37165 100644 --- a/tests/unit/user-token-limit-enforcer.test.js +++ b/tests/unit/user-token-limit-enforcer.test.js @@ -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) diff --git a/tests/unit/user-token-limits-repo.test.js b/tests/unit/user-token-limits-repo.test.js index 7da28bd3..0dde4bc5 100644 --- a/tests/unit/user-token-limits-repo.test.js +++ b/tests/unit/user-token-limits-repo.test.js @@ -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" }); diff --git a/tests/unit/user-token-quota.test.js b/tests/unit/user-token-quota.test.js index 22052d89..5c0f316d 100644 --- a/tests/unit/user-token-quota.test.js +++ b/tests/unit/user-token-quota.test.js @@ -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 () => {