From 8f6e57dfc6a7fa4e56fc9a776a14cf6f2dc0bcc8 Mon Sep 17 00:00:00 2001 From: Loi Phan Date: Sat, 18 Jul 2026 13:23:02 +0700 Subject: [PATCH] fix: update the quota tracker for user role --- CHANGELOG.md | 1 + open-sse/config/userTokenLimits.js | 11 ++ .../usage/components/SystemQuotaOverview.js | 64 ++++++- .../users/components/tokenLimitDisplay.js | 19 +-- src/app/api/usage/system-quota/route.js | 126 ++++++++++++-- .../api/users/[userId]/token-usage/route.js | 59 +------ src/lib/userTokenQuota.js | 73 ++++++++ src/shared/utils/tokenCount.js | 10 ++ tests/unit/system-quota-route.test.js | 157 ++++++++++++++++++ tests/unit/user-token-quota.test.js | 67 ++++++++ tests/unit/user-token-usage-route.test.js | 84 ++-------- 11 files changed, 516 insertions(+), 155 deletions(-) create mode 100644 src/lib/userTokenQuota.js create mode 100644 src/shared/utils/tokenCount.js create mode 100644 tests/unit/system-quota-route.test.js create mode 100644 tests/unit/user-token-quota.test.js diff --git a/CHANGELOG.md b/CHANGELOG.md index 59b506d6..c56524cb 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -27,6 +27,7 @@ ## Improvements - **Perf**: skip inactive background services on startup +- **User quota**: Quota Tracker reports personal Codex and Orbit Provider token budgets with session and weekly usage ## Docs - README: Persian YouTube tutorial diff --git a/open-sse/config/userTokenLimits.js b/open-sse/config/userTokenLimits.js index 14cd31d3..e2e81d99 100644 --- a/open-sse/config/userTokenLimits.js +++ b/open-sse/config/userTokenLimits.js @@ -16,4 +16,15 @@ export const USER_TOKEN_LIMIT_WINDOW_IDS = Object.freeze( Object.values(USER_TOKEN_LIMIT_WINDOWS), ); +export const USER_TOKEN_LIMIT_WINDOW_CONFIG = Object.freeze({ + [USER_TOKEN_LIMIT_WINDOWS.SESSION]: Object.freeze({ + name: "Session", + description: "Rolling 5 hours", + }), + [USER_TOKEN_LIMIT_WINDOWS.WEEKLY]: Object.freeze({ + name: "Weekly", + description: "Resets Monday, Vietnam time", + }), +}); + export const USER_TOKEN_LIMIT_SESSION_MS = 5 * 60 * 60 * 1000; \ No newline at end of file diff --git a/src/app/(dashboard)/dashboard/usage/components/SystemQuotaOverview.js b/src/app/(dashboard)/dashboard/usage/components/SystemQuotaOverview.js index 462cef8e..84c4f48b 100644 --- a/src/app/(dashboard)/dashboard/usage/components/SystemQuotaOverview.js +++ b/src/app/(dashboard)/dashboard/usage/components/SystemQuotaOverview.js @@ -6,6 +6,8 @@ import { Button, Card, CardSkeleton } from "@/shared/components"; import ProviderIcon from "@/shared/components/ProviderIcon"; import { 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"; function getProviderInfo(providerId) { return AI_PROVIDERS[providerId] || { @@ -63,10 +65,58 @@ 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 ? ( + + Unlimited + + ) : ( + + + )} +
    +
    + {formatTokenCount(quota.used)} used + {isUnlimited ? null : {formatTokenCount(quota.limit)} token limit} +
    + {!isUnlimited && ( +
    +
    +
    + )} +
    • + ); +} + function ProviderQuotaListItem({ provider }) { const providerInfo = getProviderInfo(provider.provider); const providerName = providerInfo.name || provider.provider; const hasQuotaData = provider.quotas.length > 0; + const isPersonalTokenBudget = provider.quotaSource === "user-token-limit"; return (
    @@ -87,16 +137,18 @@ function ProviderQuotaListItem({ provider }) {

    {providerName}

    -

    - {provider.accountCount} {provider.accountCount === 1 ? "account" : "accounts"} connected -

    + {isPersonalTokenBudget ?

    Personal token budget

    : null}
    {hasQuotaData ? ( ) : (
    @@ -106,9 +158,9 @@ function ProviderQuotaListItem({ provider }) {
    )} - {provider.failedAccountCount > 0 && hasQuotaData && ( + {provider.hasFailedQuotaChecks && hasQuotaData && (

    - Some account quota checks could not be completed. + Some quota checks could not be completed.

    )}
    diff --git a/src/app/(dashboard)/dashboard/users/components/tokenLimitDisplay.js b/src/app/(dashboard)/dashboard/users/components/tokenLimitDisplay.js index 8f2d1ec8..cd57aca7 100644 --- a/src/app/(dashboard)/dashboard/users/components/tokenLimitDisplay.js +++ b/src/app/(dashboard)/dashboard/users/components/tokenLimitDisplay.js @@ -1,13 +1,9 @@ import { USER_TOKEN_LIMIT_PROVIDERS, + USER_TOKEN_LIMIT_WINDOW_CONFIG, USER_TOKEN_LIMIT_WINDOWS, } from "open-sse/config/userTokenLimits.js"; - -const TOKEN_FORMATTER = new Intl.NumberFormat("en-US", { maximumFractionDigits: 0 }); -const COMPACT_TOKEN_FORMATTER = new Intl.NumberFormat("en-US", { - notation: "compact", - maximumFractionDigits: 1, -}); +import { formatTokenCount } from "@/shared/utils/tokenCount.js"; export const TOKEN_LIMIT_PROVIDER_OPTIONS = Object.freeze([ { @@ -29,13 +25,11 @@ export const TOKEN_LIMIT_PROVIDER_OPTIONS = Object.freeze([ export const TOKEN_LIMIT_WINDOW_OPTIONS = Object.freeze([ { id: USER_TOKEN_LIMIT_WINDOWS.SESSION, - name: "Session", - description: "Rolling 5 hours", + ...USER_TOKEN_LIMIT_WINDOW_CONFIG[USER_TOKEN_LIMIT_WINDOWS.SESSION], }, { id: USER_TOKEN_LIMIT_WINDOWS.WEEKLY, - name: "Weekly", - description: "Resets Monday, Vietnam time", + ...USER_TOKEN_LIMIT_WINDOW_CONFIG[USER_TOKEN_LIMIT_WINDOWS.WEEKLY], }, ]); @@ -70,7 +64,4 @@ export function getProviderRemainingPercentage(providerUsage) { return Math.min(...activeWindows.map((windowUsage) => windowUsage.remainingPercentage)); } -export function formatTokenCount(value, compact = false) { - const amount = Math.max(0, Number(value) || 0); - return (compact ? COMPACT_TOKEN_FORMATTER : TOKEN_FORMATTER).format(amount); -} +export { formatTokenCount }; diff --git a/src/app/api/usage/system-quota/route.js b/src/app/api/usage/system-quota/route.js index f28184cb..628506da 100644 --- a/src/app/api/usage/system-quota/route.js +++ b/src/app/api/usage/system-quota/route.js @@ -7,6 +7,12 @@ import { resolveConnectionProxyConfig } from "@/lib/network/connectionProxy"; import { USAGE_APIKEY_PROVIDERS, USAGE_SUPPORTED_PROVIDERS } from "@/shared/constants/providers"; import { refreshAndUpdateCredentials } from "@/app/api/usage/[connectionId]/route"; import { getUsageForProvider } from "open-sse/services/usage.js"; +import { getUserTokenQuota } from "@/lib/userTokenQuota.js"; +import { + USER_TOKEN_LIMIT_PROVIDER_IDS, + USER_TOKEN_LIMIT_WINDOW_CONFIG, + USER_TOKEN_LIMIT_WINDOW_IDS, +} from "open-sse/config/userTokenLimits.js"; import { getRemainingPercentage, parseQuotaData, @@ -15,7 +21,12 @@ import { export const dynamic = "force-dynamic"; const CACHE_TTL_MS = 60 * 1000; -let cachedSystemQuota = null; +const cachedSystemQuotas = new Map(); +const userTokenQuotaProviderSet = new Set(USER_TOKEN_LIMIT_PROVIDER_IDS); +const CACHE_KEYS = Object.freeze({ + ALL_PROVIDERS: "all-providers", + NON_TOKEN_PROVIDERS: "non-token-providers", +}); function isUsageEligible(connection) { const isApiKey = connection.authType === "apikey" || connection.authType === "api_key"; @@ -60,18 +71,89 @@ function sanitizeProviderError(error) { return "Quota data is temporarily unavailable for this provider."; } -function hideQuotaResetDetails(data, user) { +function sanitizeQuotaForUser(data, user) { if (user.role === "admin") return data; return { ...data, - providers: data.providers.map((provider) => ({ + providers: data.providers.map(({ + accountCount: _accountCount, + quotaAccountCount: _quotaAccountCount, + failedAccountCount, + ...provider + }) => ({ ...provider, + hasFailedQuotaChecks: failedAccountCount > 0, quotas: provider.quotas.map(({ resetAt: _resetAt, recurring: _recurring, ...quota }) => quota), })), }; } +function getCacheKey(user) { + return user.role === "admin" ? CACHE_KEYS.ALL_PROVIDERS : CACHE_KEYS.NON_TOKEN_PROVIDERS; +} + +function getUpstreamConnections(connections, user) { + const eligibleConnections = connections.filter(isUsageEligible); + if (user.role === "admin") return eligibleConnections; + + return eligibleConnections.filter((connection) => ( + !userTokenQuotaProviderSet.has(connection.provider) + )); +} + +function getActiveTokenQuotaConnectionCounts(connections) { + const counts = Object.fromEntries(USER_TOKEN_LIMIT_PROVIDER_IDS.map((provider) => [provider, 0])); + + for (const connection of connections) { + if (!connection.isActive || !isUsageEligible(connection)) continue; + if (!userTokenQuotaProviderSet.has(connection.provider)) continue; + counts[connection.provider] += 1; + } + + return counts; +} + +function buildUserTokenQuotaProviders(tokenQuota, activeConnectionCounts) { + return USER_TOKEN_LIMIT_PROVIDER_IDS.flatMap((provider) => { + const accountCount = activeConnectionCounts[provider] || 0; + if (accountCount === 0) return []; + + return [{ + provider, + accountCount, + quotaAccountCount: 1, + failedAccountCount: 0, + errorMessage: null, + quotaSource: "user-token-limit", + quotas: USER_TOKEN_LIMIT_WINDOW_IDS.map((windowType) => { + const quota = tokenQuota[provider]?.[windowType]; + return { + name: USER_TOKEN_LIMIT_WINDOW_CONFIG[windowType].name, + windowType, + tokenBudget: true, + limit: quota?.limit || 0, + used: quota?.used || 0, + remaining: quota?.remaining ?? null, + remainingPercentage: quota?.remainingPercentage ?? null, + isUnlimited: quota?.isUnlimited === true, + windowStart: quota?.windowStart || null, + }; + }), + }]; + }); +} + +function overlayUserTokenQuota(data, tokenQuota, activeConnectionCounts) { + const personalProviders = buildUserTokenQuotaProviders(tokenQuota, activeConnectionCounts); + + return { + ...data, + providers: [...data.providers, ...personalProviders] + .sort((a, b) => a.provider.localeCompare(b.provider)), + }; +} + function buildSystemQuotaResponse(connections, results) { const providerGroups = new Map(); @@ -186,15 +268,29 @@ export async function GET(request) { const { searchParams } = new URL(request.url); const forceRefresh = searchParams.get("refresh") === "true"; + const cacheKey = getCacheKey(user); + const cachedSystemQuota = cachedSystemQuotas.get(cacheKey); const cacheIsFresh = cachedSystemQuota && Date.now() - cachedSystemQuota.cachedAt < CACHE_TTL_MS; if (!forceRefresh && cacheIsFresh) { - return Response.json({ ...hideQuotaResetDetails(cachedSystemQuota.data, user), cached: true }); + if (user.role === "admin") { + return Response.json({ ...sanitizeQuotaForUser(cachedSystemQuota.data, user), cached: true }); + } + + const connections = await getProviderConnections({}); + const tokenQuota = await getUserTokenQuota(user.id); + const data = overlayUserTokenQuota( + cachedSystemQuota.data, + tokenQuota, + getActiveTokenQuotaConnectionCounts(connections), + ); + return Response.json({ ...sanitizeQuotaForUser(data, user), cached: true }); } - const connections = (await getProviderConnections({})).filter(isUsageEligible); + const connections = await getProviderConnections({}); + const upstreamConnections = getUpstreamConnections(connections, user); const results = await Promise.allSettled( - connections.map(async (connection) => ({ + upstreamConnections.map(async (connection) => ({ provider: connection.provider, quotas: await fetchConnectionQuota(connection), })), @@ -209,19 +305,27 @@ export async function GET(request) { }; } - console.warn(`[System quota] ${connections[index].provider}: ${result.reason?.message || "quota fetch failed"}`); + console.warn(`[System quota] ${upstreamConnections[index].provider}: ${result.reason?.message || "quota fetch failed"}`); return { status: result.status, - provider: connections[index].provider, + provider: upstreamConnections[index].provider, value: [], errorMessage: sanitizeProviderError(result.reason), }; }); - const data = buildSystemQuotaResponse(connections, normalizedResults); - cachedSystemQuota = { data, cachedAt: Date.now() }; + const upstreamData = buildSystemQuotaResponse(upstreamConnections, normalizedResults); + cachedSystemQuotas.set(cacheKey, { data: upstreamData, cachedAt: Date.now() }); - return Response.json({ ...hideQuotaResetDetails(data, user), cached: false }); + const data = user.role === "admin" + ? upstreamData + : overlayUserTokenQuota( + upstreamData, + await getUserTokenQuota(user.id), + getActiveTokenQuotaConnectionCounts(connections), + ); + + return Response.json({ ...sanitizeQuotaForUser(data, user), cached: false }); } catch (error) { if (error?.message === "Unauthorized") { return Response.json({ error: "Unauthorized" }, { status: 401 }); diff --git a/src/app/api/users/[userId]/token-usage/route.js b/src/app/api/users/[userId]/token-usage/route.js index 213772b0..4b1336a2 100644 --- a/src/app/api/users/[userId]/token-usage/route.js +++ b/src/app/api/users/[userId]/token-usage/route.js @@ -1,15 +1,9 @@ import { NextResponse } from "next/server"; import { getUserById, - getUserProviderTokenUsageSince, - getUserTokenLimits, } from "@/lib/db/index.js"; import { requireAdminUser } from "@/lib/auth/currentUser.js"; -import { getUserTokenLimitWindowStart } from "@/lib/tokenLimitEnforcer.js"; -import { - USER_TOKEN_LIMIT_PROVIDER_IDS, - USER_TOKEN_LIMIT_WINDOW_IDS, -} from "open-sse/config/userTokenLimits.js"; +import { getUserTokenQuota } from "@/lib/userTokenQuota.js"; export const dynamic = "force-dynamic"; @@ -27,25 +21,6 @@ function errorResponse(error) { return NextResponse.json({ error: message }, { status, headers: NO_STORE_HEADERS }); } -function buildWindowUsage(limit, used, windowStart) { - const normalizedLimit = Math.max(0, Number(limit) || 0); - const normalizedUsed = Math.max(0, Number(used) || 0); - const remaining = normalizedLimit > 0 - ? Math.max(0, normalizedLimit - normalizedUsed) - : null; - const remainingPercentage = normalizedLimit > 0 - ? Math.round((remaining / normalizedLimit) * 100) - : null; - - return { - limit: normalizedLimit, - used: normalizedUsed, - remaining, - remainingPercentage, - windowStart: windowStart.toISOString(), - }; -} - export async function GET(_request, { params }) { try { await requireAdminUser(); @@ -56,39 +31,11 @@ export async function GET(_request, { params }) { if (user.role !== "user") throw new Error("Token usage only applies to user accounts"); const now = new Date(); - const limits = await getUserTokenLimits(user.id); - const windows = Object.fromEntries(USER_TOKEN_LIMIT_WINDOW_IDS.map((windowType) => [ - windowType, - getUserTokenLimitWindowStart(windowType, now), - ])); - - const usageEntries = await Promise.all( - USER_TOKEN_LIMIT_PROVIDER_IDS.flatMap((provider) => ( - USER_TOKEN_LIMIT_WINDOW_IDS.map(async (windowType) => { - const used = await getUserProviderTokenUsageSince( - user.id, - provider, - windows[windowType], - ); - return [provider, windowType, used]; - }) - )), - ); - - const usageByProvider = Object.fromEntries( - USER_TOKEN_LIMIT_PROVIDER_IDS.map((provider) => [provider, {}]), - ); - for (const [provider, windowType, used] of usageEntries) { - usageByProvider[provider][windowType] = buildWindowUsage( - limits[provider]?.[windowType], - used, - windows[windowType], - ); - } + const providers = await getUserTokenQuota(user.id, now); return NextResponse.json({ userId: user.id, - providers: usageByProvider, + providers, updatedAt: now.toISOString(), }, { headers: NO_STORE_HEADERS }); } catch (error) { diff --git a/src/lib/userTokenQuota.js b/src/lib/userTokenQuota.js new file mode 100644 index 00000000..a454f4f5 --- /dev/null +++ b/src/lib/userTokenQuota.js @@ -0,0 +1,73 @@ +import { + getUserProviderTokenUsageSince, + getUserTokenLimits, +} from "@/lib/db/index.js"; +import { getUserTokenLimitWindowStart } from "@/lib/tokenLimitEnforcer.js"; +import { + USER_TOKEN_LIMIT_PROVIDER_IDS, + USER_TOKEN_LIMIT_WINDOW_IDS, +} from "open-sse/config/userTokenLimits.js"; + +function normalizeNonNegativeNumber(value) { + return Math.max(0, Number(value) || 0); +} + +export function buildUserTokenQuotaWindow(limit, used, windowStart) { + const normalizedLimit = normalizeNonNegativeNumber(limit); + const normalizedUsed = normalizeNonNegativeNumber(used); + const isUnlimited = normalizedLimit === 0; + const remaining = isUnlimited + ? null + : Math.max(0, normalizedLimit - normalizedUsed); + const remainingPercentage = isUnlimited + ? null + : Math.round((remaining / normalizedLimit) * 100); + + return { + limit: normalizedLimit, + used: normalizedUsed, + remaining, + remainingPercentage, + isUnlimited, + windowStart: windowStart.toISOString(), + }; +} + +/** + * Return the configured token-budget usage for a dashboard user. + * Limits of zero intentionally remain unlimited while still reporting use. + */ +export async function getUserTokenQuota(userId, now = new Date()) { + if (!userId) { + 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 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 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], + ); + } + + return providers; +} diff --git a/src/shared/utils/tokenCount.js b/src/shared/utils/tokenCount.js new file mode 100644 index 00000000..5534b1ed --- /dev/null +++ b/src/shared/utils/tokenCount.js @@ -0,0 +1,10 @@ +const TOKEN_FORMATTER = new Intl.NumberFormat("en-US", { maximumFractionDigits: 0 }); +const COMPACT_TOKEN_FORMATTER = new Intl.NumberFormat("en-US", { + notation: "compact", + maximumFractionDigits: 1, +}); + +export function formatTokenCount(value, compact = false) { + const amount = Math.max(0, Number(value) || 0); + return (compact ? COMPACT_TOKEN_FORMATTER : TOKEN_FORMATTER).format(amount); +} diff --git a/tests/unit/system-quota-route.test.js b/tests/unit/system-quota-route.test.js new file mode 100644 index 00000000..61648d41 --- /dev/null +++ b/tests/unit/system-quota-route.test.js @@ -0,0 +1,157 @@ +import { beforeEach, describe, expect, it, vi } from "vitest"; + +const getProviderConnections = vi.fn(); +const requireUsageDashboardUser = vi.fn(); +const resolveConnectionProxyConfig = vi.fn(); +const refreshAndUpdateCredentials = vi.fn(); +const getUsageForProvider = vi.fn(); +const getUserTokenQuota = vi.fn(); + +vi.mock("open-sse/index.js", () => ({})); +vi.mock("@/lib/localDb", () => ({ getProviderConnections })); +vi.mock("@/lib/auth/currentUser", () => ({ requireUsageDashboardUser })); +vi.mock("@/lib/network/connectionProxy", () => ({ resolveConnectionProxyConfig })); +vi.mock("@/app/api/usage/[connectionId]/route", () => ({ refreshAndUpdateCredentials })); +vi.mock("open-sse/services/usage.js", () => ({ getUsageForProvider })); +vi.mock("@/lib/userTokenQuota.js", () => ({ getUserTokenQuota })); + +const { GET } = await import("@/app/api/usage/system-quota/route.js"); + +const proxyConfig = { + connectionProxyEnabled: false, + connectionProxyUrl: "", + connectionNoProxy: "", + vercelRelayUrl: "", +}; + +const connections = [ + { id: "codex-1", provider: "codex", authType: "oauth", isActive: true, accessToken: "codex-token" }, + { id: "orbit-1", provider: "orbit-provider", authType: "apikey", isActive: true, apiKey: "orbit-key" }, + { id: "claude-1", provider: "claude", authType: "oauth", isActive: true, accessToken: "claude-token" }, +]; + +function quota(sessionLimit, sessionUsed, weeklyLimit, weeklyUsed) { + const buildWindow = (limit, used, windowStart) => ({ + limit, + used, + remaining: limit > 0 ? Math.max(0, limit - used) : null, + remainingPercentage: limit > 0 ? Math.round((Math.max(0, limit - used) / limit) * 100) : null, + isUnlimited: limit === 0, + windowStart, + }); + + return { + "orbit-provider": { + session: buildWindow(sessionLimit, sessionUsed, "2026-07-17T05:00:00.000Z"), + weekly: buildWindow(weeklyLimit, weeklyUsed, "2026-07-13T17:00:00.000Z"), + }, + codex: { + session: buildWindow(sessionLimit, sessionUsed, "2026-07-17T05:00:00.000Z"), + weekly: buildWindow(weeklyLimit, weeklyUsed, "2026-07-13T17:00:00.000Z"), + }, + }; +} + +function request(search = "") { + return new Request(`https://9router.local/api/usage/system-quota${search}`); +} + +function providerById(payload, provider) { + return payload.providers.find((entry) => entry.provider === provider); +} + +describe("/api/usage/system-quota personal token quota overlay", () => { + beforeEach(() => { + getProviderConnections.mockReset(); + requireUsageDashboardUser.mockReset(); + resolveConnectionProxyConfig.mockReset(); + refreshAndUpdateCredentials.mockReset(); + getUsageForProvider.mockReset(); + getUserTokenQuota.mockReset(); + + getProviderConnections.mockResolvedValue(connections); + resolveConnectionProxyConfig.mockResolvedValue(proxyConfig); + refreshAndUpdateCredentials.mockImplementation(async (connection) => ({ connection, refreshed: false })); + getUsageForProvider.mockImplementation(async (connection) => ({ + quotas: { "Primary quota": { used: connection.provider === "claude" ? 30 : 10, total: 100 } }, + })); + }); + + it("overrides Codex and Orbit with a regular user's token budgets without fetching upstream usage", async () => { + requireUsageDashboardUser.mockResolvedValue({ id: "user-1", role: "user" }); + getUserTokenQuota.mockResolvedValue(quota(100, 25, 1000, 200)); + + const response = await GET(request("?refresh=true")); + const payload = await response.json(); + + expect(response.status).toBe(200); + expect(providerById(payload, "claude").quotas).toMatchObject([ + { name: "Primary quota", remainingPercentage: 70 }, + ]); + expect(providerById(payload, "codex")).toMatchObject({ quotaSource: "user-token-limit" }); + 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 }, + ]); + expect(providerById(payload, "orbit-provider").quotas).toMatchObject([ + { name: "Session", tokenBudget: true, limit: 100, used: 25 }, + { name: "Weekly", tokenBudget: true, limit: 1000, used: 200 }, + ]); + expect(getUsageForProvider).toHaveBeenCalledTimes(1); + expect(getUsageForProvider).toHaveBeenCalledWith( + expect.objectContaining({ provider: "claude" }), + expect.any(Object), + ); + }); + + it("overlays fresh personal usage over the shared upstream cache for each regular user", async () => { + requireUsageDashboardUser.mockResolvedValueOnce({ id: "user-1", role: "user" }); + getUserTokenQuota.mockResolvedValueOnce(quota(100, 20, 0, 80)); + const first = await (await GET(request("?refresh=true"))).json(); + + requireUsageDashboardUser.mockResolvedValueOnce({ id: "user-2", role: "user" }); + getUserTokenQuota.mockResolvedValueOnce(quota(200, 30, 1000, 900)); + const second = await (await GET(request())).json(); + + expect(providerById(first, "codex").quotas).toMatchObject([ + { name: "Session", limit: 100, used: 20, remainingPercentage: 80 }, + { name: "Weekly", limit: 0, used: 80, isUnlimited: true }, + ]); + expect(providerById(second, "codex").quotas).toMatchObject([ + { name: "Session", limit: 200, used: 30, remainingPercentage: 85 }, + { name: "Weekly", limit: 1000, used: 900, remainingPercentage: 10 }, + ]); + expect(getUsageForProvider).toHaveBeenCalledTimes(1); + expect(getUserTokenQuota).toHaveBeenNthCalledWith(1, "user-1"); + expect(getUserTokenQuota).toHaveBeenNthCalledWith(2, "user-2"); + }); + + it("omits a token-budget provider when it has no active eligible connection", async () => { + requireUsageDashboardUser.mockResolvedValue({ id: "user-1", role: "user" }); + getProviderConnections.mockResolvedValue([ + { ...connections[0], isActive: false }, + connections[2], + ]); + getUserTokenQuota.mockResolvedValue(quota(100, 25, 1000, 200)); + + const payload = await (await GET(request("?refresh=true"))).json(); + + expect(providerById(payload, "codex")).toBeUndefined(); + expect(providerById(payload, "orbit-provider")).toBeUndefined(); + expect(providerById(payload, "claude")).toBeDefined(); + }); + + it("keeps upstream Codex and Orbit quota data for administrators", async () => { + requireUsageDashboardUser.mockResolvedValue({ id: "admin-1", role: "admin" }); + + const payload = await (await GET(request("?refresh=true"))).json(); + + expect(providerById(payload, "codex").quotaSource).toBeUndefined(); + expect(providerById(payload, "orbit-provider").quotaSource).toBeUndefined(); + expect(providerById(payload, "codex").quotas).toMatchObject([{ name: "Primary quota" }]); + expect(getUserTokenQuota).not.toHaveBeenCalled(); + expect(getUsageForProvider).toHaveBeenCalledTimes(3); + }); +}); diff --git a/tests/unit/user-token-quota.test.js b/tests/unit/user-token-quota.test.js new file mode 100644 index 00000000..bb12ba80 --- /dev/null +++ b/tests/unit/user-token-quota.test.js @@ -0,0 +1,67 @@ +import { beforeEach, describe, expect, it, vi } from "vitest"; + +const getUserTokenLimits = vi.fn(); +const getUserProviderTokenUsageSince = vi.fn(); +const getUserTokenLimitWindowStart = vi.fn(); + +vi.mock("@/lib/db/index.js", () => ({ + getUserTokenLimits, + getUserProviderTokenUsageSince, +})); +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"); + +function usageKey(provider, windowType) { + return `${provider}:${windowType}`; +} + +describe("user token quota snapshot", () => { + beforeEach(() => { + getUserTokenLimits.mockReset(); + getUserProviderTokenUsageSince.mockReset(); + getUserTokenLimitWindowStart.mockReset(); + + getUserTokenLimits.mockResolvedValue({ + "orbit-provider": { session: 100, weekly: 1000 }, + codex: { session: 0, weekly: 500 }, + }); + getUserTokenLimitWindowStart.mockImplementation((windowType) => ( + windowType === "session" ? sessionStart : weeklyStart + )); + const usage = new Map([ + [usageKey("orbit-provider", "session"), 25], + [usageKey("orbit-provider", "weekly"), 1200], + [usageKey("codex", "session"), 12], + [usageKey("codex", "weekly"), 400], + ]); + getUserProviderTokenUsageSince.mockImplementation(async (_userId, provider, since) => { + const windowType = since.getTime() === sessionStart.getTime() ? "session" : "weekly"; + return usage.get(usageKey(provider, windowType)); + }); + }); + + it("calculates both provider windows and preserves zero as unlimited", async () => { + const quota = await getUserTokenQuota("user-1", new Date("2026-07-17T10:00:00.000Z")); + + expect(quota).toMatchObject({ + "orbit-provider": { + session: { limit: 100, used: 25, remaining: 75, remainingPercentage: 75, isUnlimited: false }, + weekly: { limit: 1000, used: 1200, remaining: 0, remainingPercentage: 0, isUnlimited: false }, + }, + codex: { + session: { limit: 0, used: 12, remaining: null, remainingPercentage: null, isUnlimited: true }, + weekly: { limit: 500, used: 400, remaining: 100, remainingPercentage: 20, isUnlimited: false }, + }, + }); + expect(quota.codex.session.windowStart).toBe(sessionStart.toISOString()); + expect(getUserProviderTokenUsageSince).toHaveBeenCalledTimes(4); + }); + + it("requires a user id", async () => { + await expect(getUserTokenQuota()).rejects.toThrow("User id is required"); + }); +}); \ No newline at end of file diff --git a/tests/unit/user-token-usage-route.test.js b/tests/unit/user-token-usage-route.test.js index b65f2f64..9f787834 100644 --- a/tests/unit/user-token-usage-route.test.js +++ b/tests/unit/user-token-usage-route.test.js @@ -2,9 +2,7 @@ import { beforeEach, describe, expect, it, vi } from "vitest"; const requireAdminUser = vi.fn(); const getUserById = vi.fn(); -const getUserTokenLimits = vi.fn(); -const getUserProviderTokenUsageSince = vi.fn(); -const getUserTokenLimitWindowStart = vi.fn(); +const getUserTokenQuota = vi.fn(); vi.mock("next/server", () => ({ NextResponse: { @@ -19,53 +17,36 @@ vi.mock("next/server", () => ({ vi.mock("@/lib/auth/currentUser.js", () => ({ requireAdminUser })); vi.mock("@/lib/db/index.js", () => ({ getUserById, - getUserTokenLimits, - getUserProviderTokenUsageSince, })); -vi.mock("@/lib/tokenLimitEnforcer.js", () => ({ getUserTokenLimitWindowStart })); +vi.mock("@/lib/userTokenQuota.js", () => ({ getUserTokenQuota })); const { GET } = await import("@/app/api/users/[userId]/token-usage/route.js"); const context = (userId = "user-1") => ({ params: Promise.resolve({ userId }) }); const request = new Request("https://9router.local/api/users/user-1/token-usage"); -const sessionStart = new Date("2026-07-17T05:00:00.000Z"); -const weeklyStart = new Date("2026-07-13T17:00:00.000Z"); -const limits = { - "orbit-provider": { session: 100, weekly: 1000 }, - codex: { session: 0, weekly: 500 }, +const quota = { + "orbit-provider": { + session: { limit: 100, used: 25, remaining: 75, remainingPercentage: 75, isUnlimited: false, windowStart: "2026-07-17T05:00:00.000Z" }, + weekly: { limit: 1000, used: 1200, remaining: 0, remainingPercentage: 0, isUnlimited: false, windowStart: "2026-07-13T17:00:00.000Z" }, + }, + codex: { + session: { limit: 0, used: 12, remaining: null, remainingPercentage: null, isUnlimited: true, windowStart: "2026-07-17T05:00:00.000Z" }, + weekly: { limit: 500, used: 400, remaining: 100, remainingPercentage: 20, isUnlimited: false, windowStart: "2026-07-13T17:00:00.000Z" }, + }, }; -function usageKey(provider, since) { - return `${provider}|${since.toISOString()}`; -} - describe("/api/users/[userId]/token-usage", () => { beforeEach(() => { requireAdminUser.mockReset(); getUserById.mockReset(); - getUserTokenLimits.mockReset(); - getUserProviderTokenUsageSince.mockReset(); - getUserTokenLimitWindowStart.mockReset(); + getUserTokenQuota.mockReset(); requireAdminUser.mockResolvedValue({ id: "admin-1", role: "admin" }); getUserById.mockResolvedValue({ id: "user-1", role: "user", isActive: true }); - getUserTokenLimits.mockResolvedValue(limits); - getUserTokenLimitWindowStart.mockImplementation((windowType) => ( - windowType === "session" ? sessionStart : weeklyStart - )); - - const usage = new Map([ - [usageKey("orbit-provider", sessionStart), 25], - [usageKey("orbit-provider", weeklyStart), 1200], - [usageKey("codex", sessionStart), 12], - [usageKey("codex", weeklyStart), 400], - ]); - getUserProviderTokenUsageSince.mockImplementation(async (_userId, provider, since) => ( - usage.get(usageKey(provider, since)) || 0 - )); + getUserTokenQuota.mockResolvedValue(quota); }); - it("returns usage and remaining headroom for every provider window", async () => { + it("returns the shared usage and remaining headroom snapshot", async () => { const response = await GET(request, context()); const payload = await response.json(); @@ -73,43 +54,10 @@ describe("/api/users/[userId]/token-usage", () => { expect(response.headers.get("Cache-Control")).toBe("no-store"); expect(payload).toMatchObject({ userId: "user-1", - providers: { - "orbit-provider": { - session: { - limit: 100, - used: 25, - remaining: 75, - remainingPercentage: 75, - windowStart: sessionStart.toISOString(), - }, - weekly: { - limit: 1000, - used: 1200, - remaining: 0, - remainingPercentage: 0, - windowStart: weeklyStart.toISOString(), - }, - }, - codex: { - session: { - limit: 0, - used: 12, - remaining: null, - remainingPercentage: null, - windowStart: sessionStart.toISOString(), - }, - weekly: { - limit: 500, - used: 400, - remaining: 100, - remainingPercentage: 20, - windowStart: weeklyStart.toISOString(), - }, - }, - }, + providers: quota, }); expect(payload.updatedAt).toEqual(expect.any(String)); - expect(getUserProviderTokenUsageSince).toHaveBeenCalledTimes(4); + expect(getUserTokenQuota).toHaveBeenCalledWith("user-1", expect.any(Date)); }); it("requires an administrator and an existing regular user", async () => {