nezumi/9router
1
0
Fork 0
mirror of https://github.com/Nezumi-2711/9router.git synced 2026-09-22 13:38:31 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: logic reset quota by session

Browse Source
This commit is contained in:
nezumi
2026-07-20 16:18:00 +07:00
parent 0e8a21fdbe
commit 580bf9d9ae
16 changed files with 340 additions and 87 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/unit/user-token-limit-enforcer.test.js
+27
View File
@@ -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)
tests/unit/user-token-limits-repo.test.js
+25
View File
@@ -67,6 +67,31 @@ describe("user token limit repository", () => {
)).resolves.toBe(50);
});
it("keeps a fixed session anchor until its five-hour boundary passes", async () => {
const db = await import("@/lib/db/index.js");
const user = await db.createUser({ username: "session-user", password: "password", role: "user" });
await expect(db.ensureUserTokenQuotaSession(
user.id,
"codex",
"2026-07-17T06:00:00.000Z",
)).resolves.toBe("2026-07-17T06:00:00.000Z");
await expect(db.ensureUserTokenQuotaSession(
user.id,
"codex",
"2026-07-17T08:00:00.000Z",
)).resolves.toBe("2026-07-17T06:00:00.000Z");
await expect(db.ensureUserTokenQuotaSession(
user.id,
"codex",
"2026-07-17T11:01:00.000Z",
)).resolves.toBe("2026-07-17T11:01:00.000Z");
await expect(db.getUserTokenQuotaSession(user.id, "codex"))
.resolves.toBe("2026-07-17T11:01:00.000Z");
});
it("rejects negative and non-integer limits without changing stored values", async () => {
const db = await import("@/lib/db/index.js");
const user = await db.createUser({ username: "invalid-limit", password: "password", role: "user" });
tests/unit/user-token-quota.test.js
+36 -12
View File
@@ -3,19 +3,21 @@ import { beforeEach, describe, expect, it, vi } from "vitest";
const getUserTokenLimits = vi.fn();
const getUserProviderTokenUsageSince = vi.fn();
const getUserProviderEarliestTokenUsageSince = vi.fn();
const getUserTokenLimitWindowStart = vi.fn();
const getUserTokenQuotaSession = vi.fn();
const ensureUserTokenQuotaSession = vi.fn();
vi.mock("@/lib/db/index.js", () => ({
getUserTokenLimits,
getUserProviderTokenUsageSince,
getUserProviderEarliestTokenUsageSince,
getUserTokenQuotaSession,
ensureUserTokenQuotaSession,
}));
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");
const sessionStart = new Date("2026-07-17T06:00:00.000Z");
const weeklyStart = new Date("2026-07-12T17:00:00.000Z");
function usageKey(provider, windowType) {
return `${provider}:${windowType}`;
@@ -26,15 +28,14 @@ describe("user token quota snapshot", () => {
getUserTokenLimits.mockReset();
getUserProviderTokenUsageSince.mockReset();
getUserProviderEarliestTokenUsageSince.mockReset();
getUserTokenLimitWindowStart.mockReset();
getUserTokenQuotaSession.mockReset();
ensureUserTokenQuotaSession.mockReset();
getUserTokenLimits.mockResolvedValue({
"orbit-provider": { session: 100, weekly: 1000 },
codex: { session: 0, weekly: 500 },
});
getUserTokenLimitWindowStart.mockImplementation((windowType) => (
windowType === "session" ? sessionStart : weeklyStart
));
getUserTokenQuotaSession.mockResolvedValue(sessionStart.toISOString());
const usage = new Map([
[usageKey("orbit-provider", "session"), 25],
[usageKey("orbit-provider", "weekly"), 1200],
@@ -62,21 +63,44 @@ describe("user token quota snapshot", () => {
},
});
expect(quota.codex.session.windowStart).toBe(sessionStart.toISOString());
expect(quota.codex.session.resetAt).toBeNull();
expect(quota.codex.weekly.resetAt).toBe("2026-07-20T17:00:00.000Z");
expect(quota.codex.session.resetAt).toBe("2026-07-17T11:00:00.000Z");
expect(quota.codex.weekly.resetAt).toBe("2026-07-19T17:00:00.000Z");
expect(getUserProviderTokenUsageSince).toHaveBeenCalledTimes(4);
expect(getUserProviderEarliestTokenUsageSince).toHaveBeenCalledTimes(2);
expect(getUserProviderEarliestTokenUsageSince).not.toHaveBeenCalled();
});
it("reports when the next tokens leave a rolling session window", async () => {
it("resets an expired session to its full budget until another request starts one", async () => {
getUserTokenQuotaSession.mockResolvedValue("2026-07-17T05:00:00.000Z");
const quota = await getUserTokenQuota("user-1", new Date("2026-07-17T10:10:00.000Z"));
expect(quota["orbit-provider"].session).toMatchObject({
used: 0,
remaining: 100,
remainingPercentage: 100,
windowStart: null,
resetAt: null,
});
expect(quota.codex.session).toMatchObject({ used: 0, remaining: null, windowStart: null, resetAt: null });
expect(getUserProviderEarliestTokenUsageSince).not.toHaveBeenCalled();
});
it("seeds an active fixed session from legacy usage exactly once", async () => {
getUserTokenQuotaSession.mockResolvedValue(null);
getUserProviderEarliestTokenUsageSince.mockImplementation(async (_userId, provider) => (
provider === "orbit-provider" ? "2026-07-17T06:30:00.000Z" : null
));
ensureUserTokenQuotaSession.mockResolvedValue("2026-07-17T06:30:00.000Z");
const quota = await getUserTokenQuota("user-1", new Date("2026-07-17T10:00:00.000Z"));
expect(quota["orbit-provider"].session.resetAt).toBe("2026-07-17T11:30:00.000Z");
expect(quota.codex.session.resetAt).toBeNull();
expect(ensureUserTokenQuotaSession).toHaveBeenCalledWith(
"user-1",
"orbit-provider",
new Date("2026-07-17T06:30:00.000Z"),
);
});
it("requires a user id", async () => {