fix: update the quota tracker for user role

This commit is contained in:
2026-07-18 13:23:02 +07:00
parent c77127aaf3
commit 8f6e57dfc6
11 changed files with 516 additions and 155 deletions
+1
View File
@@ -27,6 +27,7 @@
## Improvements ## Improvements
- **Perf**: skip inactive background services on startup - **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 ## Docs
- README: Persian YouTube tutorial - README: Persian YouTube tutorial
+11
View File
@@ -16,4 +16,15 @@ export const USER_TOKEN_LIMIT_WINDOW_IDS = Object.freeze(
Object.values(USER_TOKEN_LIMIT_WINDOWS), 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; export const USER_TOKEN_LIMIT_SESSION_MS = 5 * 60 * 60 * 1000;
@@ -6,6 +6,8 @@ import { Button, Card, CardSkeleton } from "@/shared/components";
import ProviderIcon from "@/shared/components/ProviderIcon"; import ProviderIcon from "@/shared/components/ProviderIcon";
import { REFRESH_INTERVAL_MS } from "./ProviderLimits/utils"; import { REFRESH_INTERVAL_MS } from "./ProviderLimits/utils";
import { formatVietnamTime } from "@/shared/utils/dateTime"; 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) { function getProviderInfo(providerId) {
return AI_PROVIDERS[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 (
<li className="py-3 first:pt-0 last:pb-0">
<div className="flex items-start justify-between gap-3">
<div className="min-w-0">
<span className="block truncate text-sm font-medium text-text-main">{quota.name}</span>
<span className="mt-0.5 block text-xs text-text-muted">{description}</span>
</div>
{isUnlimited ? (
<span className="shrink-0 rounded-full border border-border-subtle bg-surface px-2 py-1 text-[11px] font-medium text-text-muted">
Unlimited
</span>
) : (
<span className={`inline-flex shrink-0 items-center gap-2 text-sm font-semibold tabular-nums ${tone.text}`}>
<span className={`h-2 w-2 rounded-full ${tone.dot}`} aria-hidden="true" />
{quota.remainingPercentage}%
</span>
)}
</div>
<div className="mt-2 flex items-center justify-between gap-3 text-xs tabular-nums text-text-muted">
<span>{formatTokenCount(quota.used)} used</span>
{isUnlimited ? null : <span>{formatTokenCount(quota.limit)} token limit</span>}
</div>
{!isUnlimited && (
<div
className="mt-2 h-1.5 overflow-hidden rounded-full bg-surface-2"
role="progressbar"
aria-label={`${quota.name} personal token quota remaining`}
aria-valuemin={0}
aria-valuemax={100}
aria-valuenow={quota.remainingPercentage}
>
<div
className={`h-full rounded-full transition-[width] duration-300 ${tone.bar}`}
style={{ width: `${Math.min(Math.max(quota.remainingPercentage, 0), 100)}%` }}
/>
</div>
)}
</li>
);
}
function ProviderQuotaListItem({ provider }) { function ProviderQuotaListItem({ provider }) {
const providerInfo = getProviderInfo(provider.provider); const providerInfo = getProviderInfo(provider.provider);
const providerName = providerInfo.name || provider.provider; const providerName = providerInfo.name || provider.provider;
const hasQuotaData = provider.quotas.length > 0; const hasQuotaData = provider.quotas.length > 0;
const isPersonalTokenBudget = provider.quotaSource === "user-token-limit";
return ( return (
<article className="flex flex-col gap-5 px-4 py-5 sm:flex-row sm:gap-8 sm:px-6 sm:py-6"> <article className="flex flex-col gap-5 px-4 py-5 sm:flex-row sm:gap-8 sm:px-6 sm:py-6">
@@ -87,16 +137,18 @@ function ProviderQuotaListItem({ provider }) {
</div> </div>
<div className="min-w-0"> <div className="min-w-0">
<h2 className="truncate font-semibold text-text-main">{providerName}</h2> <h2 className="truncate font-semibold text-text-main">{providerName}</h2>
<p className="text-sm text-text-muted"> {isPersonalTokenBudget ? <p className="text-sm text-text-muted">Personal token budget</p> : null}
{provider.accountCount} {provider.accountCount === 1 ? "account" : "accounts"} connected
</p>
</div> </div>
</div> </div>
</header> </header>
{hasQuotaData ? ( {hasQuotaData ? (
<ul className="min-w-0 flex-1 divide-y divide-border-subtle"> <ul className="min-w-0 flex-1 divide-y divide-border-subtle">
{provider.quotas.map((quota) => <QuotaListRow key={quota.name} quota={quota} />)} {provider.quotas.map((quota) => (
quota.tokenBudget
? <TokenQuotaListRow key={quota.name} quota={quota} />
: <QuotaListRow key={quota.name} quota={quota} />
))}
</ul> </ul>
) : ( ) : (
<div className="flex-1 rounded-lg border border-dashed border-border-subtle bg-bg px-4 py-4 text-sm text-text-muted"> <div className="flex-1 rounded-lg border border-dashed border-border-subtle bg-bg px-4 py-4 text-sm text-text-muted">
@@ -106,9 +158,9 @@ function ProviderQuotaListItem({ provider }) {
</div> </div>
)} )}
{provider.failedAccountCount > 0 && hasQuotaData && ( {provider.hasFailedQuotaChecks && hasQuotaData && (
<p className="self-end text-xs text-text-muted sm:max-w-44"> <p className="self-end text-xs text-text-muted sm:max-w-44">
Some account quota checks could not be completed. Some quota checks could not be completed.
</p> </p>
)} )}
</article> </article>
@@ -1,13 +1,9 @@
import { import {
USER_TOKEN_LIMIT_PROVIDERS, USER_TOKEN_LIMIT_PROVIDERS,
USER_TOKEN_LIMIT_WINDOW_CONFIG,
USER_TOKEN_LIMIT_WINDOWS, USER_TOKEN_LIMIT_WINDOWS,
} from "open-sse/config/userTokenLimits.js"; } from "open-sse/config/userTokenLimits.js";
import { formatTokenCount } from "@/shared/utils/tokenCount.js";
const TOKEN_FORMATTER = new Intl.NumberFormat("en-US", { maximumFractionDigits: 0 });
const COMPACT_TOKEN_FORMATTER = new Intl.NumberFormat("en-US", {
notation: "compact",
maximumFractionDigits: 1,
});
export const TOKEN_LIMIT_PROVIDER_OPTIONS = Object.freeze([ 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([ export const TOKEN_LIMIT_WINDOW_OPTIONS = Object.freeze([
{ {
id: USER_TOKEN_LIMIT_WINDOWS.SESSION, id: USER_TOKEN_LIMIT_WINDOWS.SESSION,
name: "Session", ...USER_TOKEN_LIMIT_WINDOW_CONFIG[USER_TOKEN_LIMIT_WINDOWS.SESSION],
description: "Rolling 5 hours",
}, },
{ {
id: USER_TOKEN_LIMIT_WINDOWS.WEEKLY, id: USER_TOKEN_LIMIT_WINDOWS.WEEKLY,
name: "Weekly", ...USER_TOKEN_LIMIT_WINDOW_CONFIG[USER_TOKEN_LIMIT_WINDOWS.WEEKLY],
description: "Resets Monday, Vietnam time",
}, },
]); ]);
@@ -70,7 +64,4 @@ export function getProviderRemainingPercentage(providerUsage) {
return Math.min(...activeWindows.map((windowUsage) => windowUsage.remainingPercentage)); return Math.min(...activeWindows.map((windowUsage) => windowUsage.remainingPercentage));
} }
export function formatTokenCount(value, compact = false) { export { formatTokenCount };
const amount = Math.max(0, Number(value) || 0);
return (compact ? COMPACT_TOKEN_FORMATTER : TOKEN_FORMATTER).format(amount);
}
+115 -11
View File
@@ -7,6 +7,12 @@ import { resolveConnectionProxyConfig } from "@/lib/network/connectionProxy";
import { USAGE_APIKEY_PROVIDERS, USAGE_SUPPORTED_PROVIDERS } from "@/shared/constants/providers"; import { USAGE_APIKEY_PROVIDERS, USAGE_SUPPORTED_PROVIDERS } from "@/shared/constants/providers";
import { refreshAndUpdateCredentials } from "@/app/api/usage/[connectionId]/route"; import { refreshAndUpdateCredentials } from "@/app/api/usage/[connectionId]/route";
import { getUsageForProvider } from "open-sse/services/usage.js"; 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 { import {
getRemainingPercentage, getRemainingPercentage,
parseQuotaData, parseQuotaData,
@@ -15,7 +21,12 @@ import {
export const dynamic = "force-dynamic"; export const dynamic = "force-dynamic";
const CACHE_TTL_MS = 60 * 1000; 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) { function isUsageEligible(connection) {
const isApiKey = connection.authType === "apikey" || connection.authType === "api_key"; 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."; return "Quota data is temporarily unavailable for this provider.";
} }
function hideQuotaResetDetails(data, user) { function sanitizeQuotaForUser(data, user) {
if (user.role === "admin") return data; if (user.role === "admin") return data;
return { return {
...data, ...data,
providers: data.providers.map((provider) => ({ providers: data.providers.map(({
accountCount: _accountCount,
quotaAccountCount: _quotaAccountCount,
failedAccountCount,
...provider
}) => ({
...provider, ...provider,
hasFailedQuotaChecks: failedAccountCount > 0,
quotas: provider.quotas.map(({ resetAt: _resetAt, recurring: _recurring, ...quota }) => quota), 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) { function buildSystemQuotaResponse(connections, results) {
const providerGroups = new Map(); const providerGroups = new Map();
@@ -186,15 +268,29 @@ export async function GET(request) {
const { searchParams } = new URL(request.url); const { searchParams } = new URL(request.url);
const forceRefresh = searchParams.get("refresh") === "true"; 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; const cacheIsFresh = cachedSystemQuota && Date.now() - cachedSystemQuota.cachedAt < CACHE_TTL_MS;
if (!forceRefresh && cacheIsFresh) { 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({})).filter(isUsageEligible); 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({});
const upstreamConnections = getUpstreamConnections(connections, user);
const results = await Promise.allSettled( const results = await Promise.allSettled(
connections.map(async (connection) => ({ upstreamConnections.map(async (connection) => ({
provider: connection.provider, provider: connection.provider,
quotas: await fetchConnectionQuota(connection), 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 { return {
status: result.status, status: result.status,
provider: connections[index].provider, provider: upstreamConnections[index].provider,
value: [], value: [],
errorMessage: sanitizeProviderError(result.reason), errorMessage: sanitizeProviderError(result.reason),
}; };
}); });
const data = buildSystemQuotaResponse(connections, normalizedResults); const upstreamData = buildSystemQuotaResponse(upstreamConnections, normalizedResults);
cachedSystemQuota = { data, cachedAt: Date.now() }; 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) { } catch (error) {
if (error?.message === "Unauthorized") { if (error?.message === "Unauthorized") {
return Response.json({ error: "Unauthorized" }, { status: 401 }); return Response.json({ error: "Unauthorized" }, { status: 401 });
@@ -1,15 +1,9 @@
import { NextResponse } from "next/server"; import { NextResponse } from "next/server";
import { import {
getUserById, getUserById,
getUserProviderTokenUsageSince,
getUserTokenLimits,
} from "@/lib/db/index.js"; } from "@/lib/db/index.js";
import { requireAdminUser } from "@/lib/auth/currentUser.js"; import { requireAdminUser } from "@/lib/auth/currentUser.js";
import { getUserTokenLimitWindowStart } from "@/lib/tokenLimitEnforcer.js"; import { getUserTokenQuota } from "@/lib/userTokenQuota.js";
import {
USER_TOKEN_LIMIT_PROVIDER_IDS,
USER_TOKEN_LIMIT_WINDOW_IDS,
} from "open-sse/config/userTokenLimits.js";
export const dynamic = "force-dynamic"; export const dynamic = "force-dynamic";
@@ -27,25 +21,6 @@ function errorResponse(error) {
return NextResponse.json({ error: message }, { status, headers: NO_STORE_HEADERS }); 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 }) { export async function GET(_request, { params }) {
try { try {
await requireAdminUser(); 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"); if (user.role !== "user") throw new Error("Token usage only applies to user accounts");
const now = new Date(); const now = new Date();
const limits = await getUserTokenLimits(user.id); const providers = await getUserTokenQuota(user.id, now);
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],
);
}
return NextResponse.json({ return NextResponse.json({
userId: user.id, userId: user.id,
providers: usageByProvider, providers,
updatedAt: now.toISOString(), updatedAt: now.toISOString(),
}, { headers: NO_STORE_HEADERS }); }, { headers: NO_STORE_HEADERS });
} catch (error) { } catch (error) {
+73
View File
@@ -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;
}
+10
View File
@@ -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);
}
+157
View File
@@ -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);
});
});
+67
View File
@@ -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");
});
});
+16 -68
View File
@@ -2,9 +2,7 @@ import { beforeEach, describe, expect, it, vi } from "vitest";
const requireAdminUser = vi.fn(); const requireAdminUser = vi.fn();
const getUserById = vi.fn(); const getUserById = vi.fn();
const getUserTokenLimits = vi.fn(); const getUserTokenQuota = vi.fn();
const getUserProviderTokenUsageSince = vi.fn();
const getUserTokenLimitWindowStart = vi.fn();
vi.mock("next/server", () => ({ vi.mock("next/server", () => ({
NextResponse: { NextResponse: {
@@ -19,53 +17,36 @@ vi.mock("next/server", () => ({
vi.mock("@/lib/auth/currentUser.js", () => ({ requireAdminUser })); vi.mock("@/lib/auth/currentUser.js", () => ({ requireAdminUser }));
vi.mock("@/lib/db/index.js", () => ({ vi.mock("@/lib/db/index.js", () => ({
getUserById, 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 { GET } = await import("@/app/api/users/[userId]/token-usage/route.js");
const context = (userId = "user-1") => ({ params: Promise.resolve({ userId }) }); const context = (userId = "user-1") => ({ params: Promise.resolve({ userId }) });
const request = new Request("https://9router.local/api/users/user-1/token-usage"); const request = new Request("https://9router.local/api/users/user-1/token-usage");
const sessionStart = new Date("2026-07-17T05:00:00.000Z"); const quota = {
const weeklyStart = new Date("2026-07-13T17:00:00.000Z"); "orbit-provider": {
const limits = { session: { limit: 100, used: 25, remaining: 75, remainingPercentage: 75, isUnlimited: false, windowStart: "2026-07-17T05:00:00.000Z" },
"orbit-provider": { session: 100, weekly: 1000 }, weekly: { limit: 1000, used: 1200, remaining: 0, remainingPercentage: 0, isUnlimited: false, windowStart: "2026-07-13T17:00:00.000Z" },
codex: { session: 0, weekly: 500 }, },
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", () => { describe("/api/users/[userId]/token-usage", () => {
beforeEach(() => { beforeEach(() => {
requireAdminUser.mockReset(); requireAdminUser.mockReset();
getUserById.mockReset(); getUserById.mockReset();
getUserTokenLimits.mockReset(); getUserTokenQuota.mockReset();
getUserProviderTokenUsageSince.mockReset();
getUserTokenLimitWindowStart.mockReset();
requireAdminUser.mockResolvedValue({ id: "admin-1", role: "admin" }); requireAdminUser.mockResolvedValue({ id: "admin-1", role: "admin" });
getUserById.mockResolvedValue({ id: "user-1", role: "user", isActive: true }); getUserById.mockResolvedValue({ id: "user-1", role: "user", isActive: true });
getUserTokenLimits.mockResolvedValue(limits); getUserTokenQuota.mockResolvedValue(quota);
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
));
}); });
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 response = await GET(request, context());
const payload = await response.json(); 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(response.headers.get("Cache-Control")).toBe("no-store");
expect(payload).toMatchObject({ expect(payload).toMatchObject({
userId: "user-1", userId: "user-1",
providers: { providers: quota,
"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(),
},
},
},
}); });
expect(payload.updatedAt).toEqual(expect.any(String)); 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 () => { it("requires an administrator and an existing regular user", async () => {