feat: limit quota for user

2026-09-22 13:38:31 +00:00 · 2026-07-17 18:08:50 +07:00
parent 2e595627d0
commit b700a070c0
15 changed files with 770 additions and 3 deletions
@@ -0,0 +1,93 @@
+import { beforeEach, describe, expect, it, vi } from "vitest";
+
+const getUserById = vi.fn();
+const getUserProviderTokenUsageSince = vi.fn();
+const getUserTokenLimits = vi.fn();
+
+vi.mock("@/lib/db/index.js", () => ({
+  getUserById,
+  getUserProviderTokenUsageSince,
+  getUserTokenLimits,
+}));
+
+const {
+  checkUserTokenLimit,
+  getUserTokenLimitWindowStart,
+} = await import("@/lib/tokenLimitEnforcer.js");
+
+describe("user token limit enforcement", () => {
+  beforeEach(() => {
+    getUserById.mockReset();
+    getUserProviderTokenUsageSince.mockReset();
+    getUserTokenLimits.mockReset();
+    getUserById.mockResolvedValue({ id: "user-1", role: "user", isActive: true });
+    getUserTokenLimits.mockResolvedValue({
+      "orbit-provider": { session: 100, weekly: 1000 },
+      codex: { session: 200, weekly: 2000 },
+    });
+  });
+
+  it("calculates rolling session and Monday Vietnam weekly window starts", () => {
+    const now = new Date("2026-07-17T10:00:00.000Z");
+
+    expect(getUserTokenLimitWindowStart("session", now).toISOString())
+      .toBe("2026-07-17T05:00:00.000Z");
+    expect(getUserTokenLimitWindowStart("weekly", now).toISOString())
+      .toBe("2026-07-12T17:00:00.000Z");
+  });
+
+  it("blocks when the rolling session total reaches its limit", async () => {
+    getUserProviderTokenUsageSince.mockResolvedValueOnce(100);
+
+    const result = await checkUserTokenLimit(
+      "user-1",
+      "orbit-provider",
+      new Date("2026-07-17T10:00:00.000Z"),
+    );
+
+    expect(result).toMatchObject({
+      exceeded: true,
+      provider: "orbit-provider",
+      windowType: "session",
+      limit: 100,
+      used: 100,
+    });
+    expect(getUserProviderTokenUsageSince).toHaveBeenCalledTimes(1);
+  });
+
+  it("checks weekly usage after the session window still has headroom", async () => {
+    getUserProviderTokenUsageSince
+      .mockResolvedValueOnce(80)
+      .mockResolvedValueOnce(1200);
+
+    const result = await checkUserTokenLimit(
+      "user-1",
+      "orbit-provider",
+      new Date("2026-07-17T10:00:00.000Z"),
+    );
+
+    expect(result).toMatchObject({ windowType: "weekly", limit: 1000, used: 1200 });
+    expect(getUserProviderTokenUsageSince).toHaveBeenCalledTimes(2);
+  });
+
+  it("exempts administrators, unknown users, and unsupported providers", async () => {
+    getUserById.mockResolvedValueOnce({ id: "admin-1", role: "admin", isActive: true });
+    await expect(checkUserTokenLimit("admin-1", "codex")).resolves.toBeNull();
+
+    getUserById.mockResolvedValueOnce(null);
+    await expect(checkUserTokenLimit("missing", "codex")).resolves.toBeNull();
+
+    await expect(checkUserTokenLimit("user-1", "openai")).resolves.toBeNull();
+    expect(getUserTokenLimits).not.toHaveBeenCalled();
+  });
+
+  it("treats zero limits as unlimited", async () => {
+    getUserTokenLimits.mockResolvedValue({
+      "orbit-provider": { session: 0, weekly: 0 },
+      codex: { session: 0, weekly: 0 },
+    });
+
+    await expect(checkUserTokenLimit("user-1", "codex")).resolves.toBeNull();
+    expect(getUserProviderTokenUsageSince).not.toHaveBeenCalled();
+  });
+});