fix: update the quota tracker for user role

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);
+  });
+});

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");
+  });
+});

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 () => {