fix: logic reset quota by session

This commit is contained in:
2026-07-20 16:18:00 +07:00
parent 0e8a21fdbe
commit 580bf9d9ae
16 changed files with 340 additions and 87 deletions
@@ -2,11 +2,17 @@ import { beforeEach, describe, expect, it, vi } from "vitest";
const getUserById = vi.fn();
const getUserProviderTokenUsageSince = vi.fn();
const getUserProviderEarliestTokenUsageSince = vi.fn();
const getUserTokenQuotaSession = vi.fn();
const ensureUserTokenQuotaSession = vi.fn();
const getUserTokenLimits = vi.fn();
vi.mock("@/lib/db/index.js", () => ({
getUserById,
ensureUserTokenQuotaSession,
getUserProviderEarliestTokenUsageSince,
getUserProviderTokenUsageSince,
getUserTokenQuotaSession,
getUserTokenLimits,
}));
@@ -18,13 +24,17 @@ const {
describe("user token limit enforcement", () => {
beforeEach(() => {
getUserById.mockReset();
ensureUserTokenQuotaSession.mockReset();
getUserProviderEarliestTokenUsageSince.mockReset();
getUserProviderTokenUsageSince.mockReset();
getUserTokenQuotaSession.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 },
});
getUserTokenQuotaSession.mockResolvedValue("2026-07-17T06:00:00.000Z");
});
it("calculates rolling session and Monday Vietnam weekly window starts", () => {
@@ -55,6 +65,23 @@ describe("user token limit enforcement", () => {
expect(getUserProviderTokenUsageSince).toHaveBeenCalledTimes(1);
});
it("does not block from an expired fixed session until a new request begins one", async () => {
getUserTokenQuotaSession.mockResolvedValue("2026-07-17T05:00:00.000Z");
getUserProviderTokenUsageSince.mockResolvedValueOnce(0);
await expect(checkUserTokenLimit(
"user-1",
"orbit-provider",
new Date("2026-07-17T10:10:00.000Z"),
)).resolves.toBeNull();
expect(getUserProviderTokenUsageSince).toHaveBeenCalledWith(
"user-1",
"orbit-provider",
new Date("2026-07-12T17:00:00.000Z"),
);
});
it("checks weekly usage after the session window still has headroom", async () => {
getUserProviderTokenUsageSince
.mockResolvedValueOnce(80)