From 21532e0f00b675c15ecd15d2863667df06b1ba69 Mon Sep 17 00:00:00 2001 From: Loi Phan Date: Sat, 18 Jul 2026 16:03:34 +0700 Subject: [PATCH] fix: improve the UI for set the limit token for the user --- open-sse/config/userTokenLimits.js | 3 +- .../usage/components/ProviderLimits/utils.js | 8 +- .../usage/components/SystemQuotaOverview.js | 40 ++++-- src/app/(dashboard)/dashboard/users/page.js | 135 ++++++++++++++---- src/app/api/usage/system-quota/route.js | 7 +- src/lib/db/index.js | 1 + src/lib/db/repos/userTokenLimitsRepo.js | 24 ++++ src/lib/userTokenQuota.js | 38 +++++ tests/unit/system-quota-route.test.js | 13 +- tests/unit/user-token-quota.test.js | 18 +++ 10 files changed, 239 insertions(+), 48 deletions(-) diff --git a/open-sse/config/userTokenLimits.js b/open-sse/config/userTokenLimits.js index e2e81d99..4cccd032 100644 --- a/open-sse/config/userTokenLimits.js +++ b/open-sse/config/userTokenLimits.js @@ -27,4 +27,5 @@ export const USER_TOKEN_LIMIT_WINDOW_CONFIG = Object.freeze({ }), }); -export const USER_TOKEN_LIMIT_SESSION_MS = 5 * 60 * 60 * 1000; \ No newline at end of file +export const USER_TOKEN_LIMIT_SESSION_MS = 5 * 60 * 60 * 1000; +export const USER_TOKEN_LIMIT_WEEKLY_MS = 7 * 24 * 60 * 60 * 1000; \ No newline at end of file diff --git a/src/app/(dashboard)/dashboard/usage/components/ProviderLimits/utils.js b/src/app/(dashboard)/dashboard/usage/components/ProviderLimits/utils.js index f338073f..3d072ae8 100644 --- a/src/app/(dashboard)/dashboard/usage/components/ProviderLimits/utils.js +++ b/src/app/(dashboard)/dashboard/usage/components/ProviderLimits/utils.js @@ -211,15 +211,17 @@ export function setQuotaCache(connectionId, quotaEntry) { /** * Format ISO date string to countdown format (inspired by vscode-antigravity-cockpit) * @param {string|Date} date - ISO date string or Date object + * @param {string|Date|number} now - Current time, injectable for live countdowns * @returns {string} Formatted countdown (e.g., "2d 5h 30m", "4h 40m", "15m") or "-" */ -export function formatResetTime(date) { +export function formatResetTime(date, now = new Date()) { if (!date) return "-"; try { const resetDate = typeof date === "string" ? new Date(date) : date; - const now = new Date(); - const diffMs = resetDate - now; + const currentTime = now instanceof Date ? now : new Date(now); + if (!Number.isFinite(resetDate.getTime()) || !Number.isFinite(currentTime.getTime())) return "-"; + const diffMs = resetDate - currentTime; if (diffMs <= 0) return "-"; diff --git a/src/app/(dashboard)/dashboard/usage/components/SystemQuotaOverview.js b/src/app/(dashboard)/dashboard/usage/components/SystemQuotaOverview.js index 84c4f48b..be05bb85 100644 --- a/src/app/(dashboard)/dashboard/usage/components/SystemQuotaOverview.js +++ b/src/app/(dashboard)/dashboard/usage/components/SystemQuotaOverview.js @@ -4,10 +4,10 @@ import { useCallback, useEffect, useState } from "react"; import { AI_PROVIDERS } from "@/shared/constants/providers"; import { Button, Card, CardSkeleton } from "@/shared/components"; import ProviderIcon from "@/shared/components/ProviderIcon"; -import { REFRESH_INTERVAL_MS } from "./ProviderLimits/utils"; +import { formatResetTime, REFRESH_INTERVAL_MS } from "./ProviderLimits/utils"; import { formatVietnamTime } from "@/shared/utils/dateTime"; import { formatTokenCount } from "@/shared/utils/tokenCount.js"; -import { USER_TOKEN_LIMIT_WINDOW_CONFIG } from "open-sse/config/userTokenLimits.js"; +import { USER_TOKEN_LIMIT_WINDOWS } from "open-sse/config/userTokenLimits.js"; function getProviderInfo(providerId) { return AI_PROVIDERS[providerId] || { @@ -36,6 +36,29 @@ function getQuotaTone(percentage) { return { bar: "bg-red-500", dot: "bg-red-500", text: "text-red-500" }; } +function TokenQuotaResetStatus({ quota }) { + const [now, setNow] = useState(() => new Date()); + + useEffect(() => { + if (!quota.resetAt) return undefined; + + const intervalId = window.setInterval(() => setNow(new Date()), 60 * 1000); + return () => window.clearInterval(intervalId); + }, [quota.resetAt]); + + const countdown = formatResetTime(quota.resetAt, now); + const isSession = quota.windowType === USER_TOKEN_LIMIT_WINDOWS.SESSION; + const text = countdown === "-" + ? (isSession ? "No tokens pending expiry" : "Reset time unavailable") + : (isSession ? `Next tokens restore in ${countdown}` : `Resets in ${countdown}`); + + return ( + + {text} + + ); +} + function QuotaListRow({ quota }) { const tone = getQuotaTone(quota.remainingPercentage); @@ -68,17 +91,16 @@ function QuotaListRow({ quota }) { function TokenQuotaListRow({ quota }) { const isUnlimited = quota.isUnlimited === true; const tone = isUnlimited ? null : getQuotaTone(quota.remainingPercentage); - const windowConfig = USER_TOKEN_LIMIT_WINDOW_CONFIG[quota.windowType]; - const description = windowConfig?.description || "Personal token budget"; return (
  • {quota.name} - {description} +
    - {isUnlimited ? ( +
    + {isUnlimited ? ( Unlimited @@ -88,10 +110,9 @@ function TokenQuotaListRow({ quota }) { {quota.remainingPercentage}% )} +
    -
    - {formatTokenCount(quota.used)} used - {isUnlimited ? null : {formatTokenCount(quota.limit)} token limit} +
    {!isUnlimited && (

    {providerName}

    - {isPersonalTokenBudget ?

    Personal token budget

    : null}
    diff --git a/src/app/(dashboard)/dashboard/users/page.js b/src/app/(dashboard)/dashboard/users/page.js index a334a94f..8eb492b6 100644 --- a/src/app/(dashboard)/dashboard/users/page.js +++ b/src/app/(dashboard)/dashboard/users/page.js @@ -9,7 +9,11 @@ import { formatVietnamDateTime } from "@/shared/utils/dateTime"; import { USER_TOKEN_LIMIT_WINDOWS } from "open-sse/config/userTokenLimits.js"; import QuotaCell from "./components/QuotaCell"; import TokenLimitsUsage from "./components/TokenLimitsUsage"; -import { TOKEN_LIMIT_PROVIDER_OPTIONS } from "./components/tokenLimitDisplay.js"; +import { + TOKEN_LIMIT_PROVIDER_OPTIONS, + TOKEN_LIMIT_WINDOW_OPTIONS, + formatTokenCount, +} from "./components/tokenLimitDisplay.js"; const EMPTY_FORM = { username: "", password: "", role: "user", isActive: true }; @@ -28,6 +32,81 @@ function formatDate(value) { return formatVietnamDateTime(value, { dateStyle: "medium", timeStyle: "short" }) || "—"; } +function normalizeTokenLimitInput(value) { + return value.replace(/\D/g, "").replace(/^0+(?=\d)/, ""); +} + +function getTokenLimitNumber(value) { + const parsed = Number(value); + return Number.isSafeInteger(parsed) && parsed > 0 ? parsed : 0; +} + +function TokenLimitField({ provider, windowOption, value, onChange }) { + const [isEditing, setIsEditing] = useState(false); + const [draft, setDraft] = useState(String(value ?? "")); + const limit = getTokenLimitNumber(value); + const isUnlimited = limit === 0; + const inputId = `${provider.id}-${windowOption.id}-token-limit`; + const descriptionId = `${inputId}-description`; + + const handleChange = (event) => { + const nextValue = normalizeTokenLimitInput(event.target.value); + setDraft(nextValue); + onChange(nextValue); + }; + + const handleBlur = () => { + setIsEditing(false); + if (!draft) onChange(0); + }; + + return ( +
    +
    +
    +
    + {windowOption.id === USER_TOKEN_LIMIT_WINDOWS.SESSION ? "schedule" : "date_range"} +
    +
    +

    {windowOption.name}

    +

    {windowOption.description}

    +
    +
    + + {isUnlimited ? "Unlimited" : "Limited"} + +
    + +
    +
    + +

    {isUnlimited ? "No usage cap" : `${formatTokenCount(limit)} tokens`}

    +
    +
    + { + setDraft(String(value ?? "")); + setIsEditing(true); + }} + onBlur={handleBlur} + onChange={handleChange} + aria-describedby={descriptionId} + aria-label={`${provider.name} ${windowOption.name.toLowerCase()} token limit`} + inputClassName="h-12 pr-20 font-mono text-lg font-semibold tabular-nums" + /> + tokens +
    +

    Enter 0 to leave this window unlimited.

    +
    +
    + ); +} + export default function UsersPage() { const router = useRouter(); const user = useUserStore((state) => state.user); @@ -45,6 +124,9 @@ export default function UsersPage() { const [limitsSaving, setLimitsSaving] = useState(false); const [limitsError, setLimitsError] = useState(""); const [quotaRefreshKey, setQuotaRefreshKey] = useState(0); + const activeTokenLimitCount = TOKEN_LIMIT_PROVIDER_OPTIONS.reduce((count, provider) => ( + count + TOKEN_LIMIT_WINDOW_OPTIONS.filter(({ id }) => getTokenLimitNumber(tokenLimits[provider.id]?.[id]) > 0).length + ), 0); const loadUsers = useCallback(async () => { setLoading(true); @@ -267,7 +349,7 @@ export default function UsersPage() { isOpen={!!limitEditor} onClose={() => !limitsSaving && setLimitEditor(null)} title={`Usage & limits · ${limitEditor?.username || "user"}`} - size="xl" + size="full" footer={<>} >
    @@ -297,13 +379,19 @@ export default function UsersPage() {
    Loading token limits…
    ) : (
    -
    -

    Budget settings

    -

    Set 0 to leave a provider window unlimited.

    +
    +
    +

    Budget settings

    +

    Enter an exact token budget for each usage window. Set 0 for unlimited.

    +
    +
    +

    {activeTokenLimitCount} / {TOKEN_LIMIT_PROVIDER_OPTIONS.length * TOKEN_LIMIT_WINDOW_OPTIONS.length}

    +

    active budgets

    +
    {TOKEN_LIMIT_PROVIDER_OPTIONS.map((provider) => ( -
    -
    +
    +
    {provider.icon}
    @@ -312,31 +400,16 @@ export default function UsersPage() {

    {provider.description}

    -
    -
    - - updateTokenLimit(provider.id, USER_TOKEN_LIMIT_WINDOWS.SESSION, event.target.value)} - aria-label={`${provider.name} session token limit`} +
    + {TOKEN_LIMIT_WINDOW_OPTIONS.map((windowOption) => ( + updateTokenLimit(provider.id, windowOption.id, value)} /> -
    -
    - - updateTokenLimit(provider.id, USER_TOKEN_LIMIT_WINDOWS.WEEKLY, event.target.value)} - aria-label={`${provider.name} weekly token limit`} - /> -
    + ))}
    ))} diff --git a/src/app/api/usage/system-quota/route.js b/src/app/api/usage/system-quota/route.js index 628506da..226af440 100644 --- a/src/app/api/usage/system-quota/route.js +++ b/src/app/api/usage/system-quota/route.js @@ -84,7 +84,11 @@ function sanitizeQuotaForUser(data, user) { }) => ({ ...provider, hasFailedQuotaChecks: failedAccountCount > 0, - quotas: provider.quotas.map(({ resetAt: _resetAt, recurring: _recurring, ...quota }) => quota), + quotas: provider.quotas.map(({ resetAt, recurring: _recurring, ...quota }) => ( + provider.quotaSource === "user-token-limit" + ? { ...quota, resetAt } + : quota + )), })), }; } @@ -138,6 +142,7 @@ function buildUserTokenQuotaProviders(tokenQuota, activeConnectionCounts) { remainingPercentage: quota?.remainingPercentage ?? null, isUnlimited: quota?.isUnlimited === true, windowStart: quota?.windowStart || null, + resetAt: quota?.resetAt || null, }; }), }]; diff --git a/src/lib/db/index.js b/src/lib/db/index.js index ff478c18..9ca342b4 100644 --- a/src/lib/db/index.js +++ b/src/lib/db/index.js @@ -63,6 +63,7 @@ export { export { createEmptyUserTokenLimits, getUserTokenLimits, replaceUserTokenLimits, getUserProviderTokenUsageSince, + getUserProviderEarliestTokenUsageSince, } from "./repos/userTokenLimitsRepo.js"; // Aliases (model + custom + mitm) diff --git a/src/lib/db/repos/userTokenLimitsRepo.js b/src/lib/db/repos/userTokenLimitsRepo.js index e4a8c0c6..f146afe2 100644 --- a/src/lib/db/repos/userTokenLimitsRepo.js +++ b/src/lib/db/repos/userTokenLimitsRepo.js @@ -112,4 +112,28 @@ export async function getUserProviderTokenUsageSince(userId, provider, since) { [userId, provider, since.toISOString()], ); return Math.max(0, Number(row?.totalTokens) || 0); +} + +/** + * Return the earliest request with billable tokens inside a rolling window. + * Its expiry determines when the first part of a rolling quota becomes available again. + */ +export async function getUserProviderEarliestTokenUsageSince(userId, provider, since) { + if (!userId) return null; + assertProvider(provider); + if (!(since instanceof Date) || !Number.isFinite(since.getTime())) { + throw new Error("A valid usage window start is required"); + } + + const db = await getAdapter(); + const row = db.get( + `SELECT MIN(timestamp) AS timestamp + FROM usageHistory + WHERE userId = ? + AND provider = ? + AND timestamp >= ? + AND (COALESCE(promptTokens, 0) + COALESCE(completionTokens, 0)) > 0`, + [userId, provider, since.toISOString()], + ); + return row?.timestamp || null; } \ No newline at end of file diff --git a/src/lib/userTokenQuota.js b/src/lib/userTokenQuota.js index a454f4f5..531fae15 100644 --- a/src/lib/userTokenQuota.js +++ b/src/lib/userTokenQuota.js @@ -1,4 +1,5 @@ import { + getUserProviderEarliestTokenUsageSince, getUserProviderTokenUsageSince, getUserTokenLimits, } from "@/lib/db/index.js"; @@ -6,6 +7,9 @@ import { getUserTokenLimitWindowStart } from "@/lib/tokenLimitEnforcer.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"; function normalizeNonNegativeNumber(value) { @@ -33,6 +37,19 @@ 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(); +} + /** * Return the configured token-budget usage for a dashboard user. * Limits of zero intentionally remain unlimited while still reporting use. @@ -58,6 +75,24 @@ export async function getUserTokenQuota(userId, now = new Date()) { )), ); + 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, {}]), ); @@ -67,6 +102,9 @@ export async function getUserTokenQuota(userId, now = new Date()) { used, windows[windowType], ); + providers[provider][windowType].resetAt = windowType === USER_TOKEN_LIMIT_WINDOWS.SESSION + ? sessionNextTokenRestoreAt[provider] + : weeklyResetAt; } return providers; diff --git a/tests/unit/system-quota-route.test.js b/tests/unit/system-quota-route.test.js index 61648d41..fe720160 100644 --- a/tests/unit/system-quota-route.test.js +++ b/tests/unit/system-quota-route.test.js @@ -38,6 +38,9 @@ function quota(sessionLimit, sessionUsed, weeklyLimit, weeklyUsed) { remainingPercentage: limit > 0 ? Math.round((Math.max(0, limit - used) / limit) * 100) : null, isUnlimited: limit === 0, windowStart, + resetAt: windowStart === "2026-07-17T05:00:00.000Z" + ? "2026-07-17T10:00:00.000Z" + : "2026-07-20T17:00:00.000Z", }); return { @@ -92,8 +95,14 @@ describe("/api/usage/system-quota personal token quota overlay", () => { expect(providerById(payload, "codex").accountCount).toBeUndefined(); expect(providerById(payload, "codex").quotaAccountCount).toBeUndefined(); expect(providerById(payload, "codex").quotas).toMatchObject([ - { name: "Session", tokenBudget: true, limit: 100, used: 25, remaining: 75, remainingPercentage: 75 }, - { name: "Weekly", tokenBudget: true, limit: 1000, used: 200, remaining: 800, remainingPercentage: 80 }, + { + name: "Session", tokenBudget: true, limit: 100, used: 25, + remaining: 75, remainingPercentage: 75, resetAt: "2026-07-17T10:00:00.000Z", + }, + { + name: "Weekly", tokenBudget: true, limit: 1000, used: 200, + remaining: 800, remainingPercentage: 80, resetAt: "2026-07-20T17:00:00.000Z", + }, ]); expect(providerById(payload, "orbit-provider").quotas).toMatchObject([ { name: "Session", tokenBudget: true, limit: 100, used: 25 }, diff --git a/tests/unit/user-token-quota.test.js b/tests/unit/user-token-quota.test.js index bb12ba80..22052d89 100644 --- a/tests/unit/user-token-quota.test.js +++ b/tests/unit/user-token-quota.test.js @@ -2,11 +2,13 @@ import { beforeEach, describe, expect, it, vi } from "vitest"; const getUserTokenLimits = vi.fn(); const getUserProviderTokenUsageSince = vi.fn(); +const getUserProviderEarliestTokenUsageSince = vi.fn(); const getUserTokenLimitWindowStart = vi.fn(); vi.mock("@/lib/db/index.js", () => ({ getUserTokenLimits, getUserProviderTokenUsageSince, + getUserProviderEarliestTokenUsageSince, })); vi.mock("@/lib/tokenLimitEnforcer.js", () => ({ getUserTokenLimitWindowStart })); @@ -23,6 +25,7 @@ describe("user token quota snapshot", () => { beforeEach(() => { getUserTokenLimits.mockReset(); getUserProviderTokenUsageSince.mockReset(); + getUserProviderEarliestTokenUsageSince.mockReset(); getUserTokenLimitWindowStart.mockReset(); getUserTokenLimits.mockResolvedValue({ @@ -42,6 +45,7 @@ describe("user token quota snapshot", () => { const windowType = since.getTime() === sessionStart.getTime() ? "session" : "weekly"; return usage.get(usageKey(provider, windowType)); }); + getUserProviderEarliestTokenUsageSince.mockResolvedValue(null); }); it("calculates both provider windows and preserves zero as unlimited", async () => { @@ -58,7 +62,21 @@ 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(getUserProviderTokenUsageSince).toHaveBeenCalledTimes(4); + expect(getUserProviderEarliestTokenUsageSince).toHaveBeenCalledTimes(2); + }); + + it("reports when the next tokens leave a rolling session window", async () => { + getUserProviderEarliestTokenUsageSince.mockImplementation(async (_userId, provider) => ( + provider === "orbit-provider" ? "2026-07-17T06:30:00.000Z" : null + )); + + 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(); }); it("requires a user id", async () => {