claude-code-hub/tests/unit/proxy/rate-limit-guard.test.ts

import { beforeEach, describe, expect, it, vi } from "vitest";

const generateSessionIdMock = vi.hoisted(() => vi.fn(() => "sess_generated"));

const rateLimitServiceMock = {
  checkTotalCostLimit: vi.fn(),
  checkAndTrackKeyUserSession: vi.fn(),
  checkRpmLimit: vi.fn(),
  checkCostLimitsWithLease: vi.fn(),
  checkUserDailyCost: vi.fn(),
};

vi.mock("@/lib/rate-limit", () => ({
  RateLimitService: rateLimitServiceMock,
}));

vi.mock("@/lib/session-manager", () => ({
  SessionManager: {
    generateSessionId: generateSessionIdMock,
  },
}));

vi.mock("@/lib/logger", () => ({
  logger: {
    warn: vi.fn(),
    info: vi.fn(),
    error: vi.fn(),
    debug: vi.fn(),
  },
}));

vi.mock("next-intl/server", () => ({
  getLocale: vi.fn(async () => "zh-CN"),
}));

const getErrorMessageServerMock = vi.fn(async () => "mock rate limit message");

vi.mock("@/lib/utils/error-messages", () => ({
  ERROR_CODES: {
    RATE_LIMIT_TOTAL_EXCEEDED: "RATE_LIMIT_TOTAL_EXCEEDED",
    RATE_LIMIT_CONCURRENT_SESSIONS_EXCEEDED: "RATE_LIMIT_CONCURRENT_SESSIONS_EXCEEDED",
    RATE_LIMIT_RPM_EXCEEDED: "RATE_LIMIT_RPM_EXCEEDED",
    RATE_LIMIT_DAILY_QUOTA_EXCEEDED: "RATE_LIMIT_DAILY_QUOTA_EXCEEDED",
    RATE_LIMIT_5H_EXCEEDED: "RATE_LIMIT_5H_EXCEEDED",
    RATE_LIMIT_WEEKLY_EXCEEDED: "RATE_LIMIT_WEEKLY_EXCEEDED",
    RATE_LIMIT_MONTHLY_EXCEEDED: "RATE_LIMIT_MONTHLY_EXCEEDED",
  },
  getErrorMessageServer: getErrorMessageServerMock,
}));

describe("ProxyRateLimitGuard - key daily limit enforcement", () => {
  const createSession = (overrides?: {
    user?: Partial<{
      id: number;
      rpm: number | null;
      dailyQuota: number | null;
      dailyResetMode: "fixed" | "rolling";
      dailyResetTime: string;
      limit5hUsd: number | null;
      limitWeeklyUsd: number | null;
      limitMonthlyUsd: number | null;
      limitTotalUsd: number | null;
      limitConcurrentSessions: number | null;
    }>;
    key?: Partial<{
      id: number;
      key: string;
      limit5hUsd: number | null;
      limitDailyUsd: number | null;
      dailyResetMode: "fixed" | "rolling";
      dailyResetTime: string;
      limitWeeklyUsd: number | null;
      limitMonthlyUsd: number | null;
      limitTotalUsd: number | null;
      limitConcurrentSessions: number;
    }>;
  }) => {
    const session = {
      sessionId: "sess_test",
      authState: {
        user: {
          id: 1,
          rpm: null,
          dailyQuota: null,
          dailyResetMode: "fixed",
          dailyResetTime: "00:00",
          limit5hUsd: null,
          limitWeeklyUsd: null,
          limitMonthlyUsd: null,
          limitTotalUsd: null,
          limitConcurrentSessions: null,
          ...overrides?.user,
        },
        key: {
          id: 2,
          key: "k_test",
          limit5hUsd: null,
          limitDailyUsd: null,
          dailyResetMode: "fixed",
          dailyResetTime: "00:00",
          limitWeeklyUsd: null,
          limitMonthlyUsd: null,
          limitTotalUsd: null,
          limitConcurrentSessions: 0,
          ...overrides?.key,
        },
      },
    } as any;

    session.setSessionId = (id: string) => {
      session.sessionId = id;
    };

    return session;
  };

  beforeEach(() => {
    vi.clearAllMocks();
    generateSessionIdMock.mockReturnValue("sess_generated");

    rateLimitServiceMock.checkTotalCostLimit.mockResolvedValue({ allowed: true });
    rateLimitServiceMock.checkAndTrackKeyUserSession.mockResolvedValue({
      allowed: true,
      keyCount: 0,
      userCount: 0,
      trackedKey: false,
      trackedUser: false,
    });
    rateLimitServiceMock.checkRpmLimit.mockResolvedValue({ allowed: true });
    rateLimitServiceMock.checkUserDailyCost.mockResolvedValue({ allowed: true });
    rateLimitServiceMock.checkCostLimitsWithLease.mockResolvedValue({ allowed: true });
  });

  it("当用户未设置每日额度时，Key 每日额度已超限也必须拦截", async () => {
    const { ProxyRateLimitGuard } = await import("@/app/v1/_lib/proxy/rate-limit-guard");

    rateLimitServiceMock.checkCostLimitsWithLease
      .mockResolvedValueOnce({ allowed: true }) // key 5h
      .mockResolvedValueOnce({ allowed: true }) // user 5h
      .mockResolvedValueOnce({
        allowed: false,
        reason: "Key daily cost limit reached (usage: 20.0000/10.0000)",
      }); // key daily

    const session = createSession({
      user: { dailyQuota: null },
      key: { limitDailyUsd: 10 },
    });

    await expect(ProxyRateLimitGuard.ensure(session)).rejects.toMatchObject({
      name: "RateLimitError",
      limitType: "daily_quota",
      currentUsage: 20,
      limitValue: 10,
    });

    expect(rateLimitServiceMock.checkUserDailyCost).not.toHaveBeenCalled();

    expect(rateLimitServiceMock.checkCostLimitsWithLease).toHaveBeenCalledWith(2, "key", {
      limit_5h_usd: null,
      limit_daily_usd: 10,
      daily_reset_mode: "fixed",
      daily_reset_time: "00:00",
      limit_weekly_usd: null,
      limit_monthly_usd: null,
      cost_reset_at: null,
    });
  });

  it("当 Key 每日额度超限时，应在用户每日检查之前直接拦截（Key 优先）", async () => {
    const { ProxyRateLimitGuard } = await import("@/app/v1/_lib/proxy/rate-limit-guard");

    rateLimitServiceMock.checkCostLimitsWithLease
      .mockResolvedValueOnce({ allowed: true }) // key 5h
      .mockResolvedValueOnce({ allowed: true }) // user 5h
      .mockResolvedValueOnce({
        allowed: false,
        reason: "Key daily cost limit reached (usage: 20.0000/10.0000)",
      }); // key daily

    const session = createSession({
      user: { dailyQuota: 999 },
      key: { limitDailyUsd: 10 },
    });

    await expect(ProxyRateLimitGuard.ensure(session)).rejects.toMatchObject({
      name: "RateLimitError",
      limitType: "daily_quota",
    });

    expect(rateLimitServiceMock.checkUserDailyCost).not.toHaveBeenCalled();
  });

  it("当 Key 未设置每日额度且用户每日额度已超限时，仍应拦截用户每日额度", async () => {
    const { ProxyRateLimitGuard } = await import("@/app/v1/_lib/proxy/rate-limit-guard");

    rateLimitServiceMock.checkCostLimitsWithLease
      .mockResolvedValueOnce({ allowed: true }) // key 5h
      .mockResolvedValueOnce({ allowed: true }) // user 5h
      .mockResolvedValueOnce({ allowed: true }) // key daily (limit null)
      .mockResolvedValueOnce({
        allowed: false,
        reason: "User daily cost limit reached (usage: 20.0000/10.0000)",
      }); // user daily

    const session = createSession({
      user: { dailyQuota: 10 },
      key: { limitDailyUsd: null },
    });

    await expect(ProxyRateLimitGuard.ensure(session)).rejects.toMatchObject({
      name: "RateLimitError",
      limitType: "daily_quota",
      currentUsage: 20,
      limitValue: 10,
    });

    // User daily 现在使用 checkCostLimitsWithLease 而不是 checkUserDailyCost
    expect(rateLimitServiceMock.checkUserDailyCost).not.toHaveBeenCalled();
    expect(rateLimitServiceMock.checkCostLimitsWithLease).toHaveBeenCalledWith(1, "user", {
      limit_5h_usd: null,
      limit_daily_usd: 10,
      daily_reset_time: "00:00",
      daily_reset_mode: "fixed",
      limit_weekly_usd: null,
      limit_monthly_usd: null,
      cost_reset_at: null,
    });
  });

  it("Key 总限额超限应拦截（usd_total）", async () => {
    const { ProxyRateLimitGuard } = await import("@/app/v1/_lib/proxy/rate-limit-guard");

    rateLimitServiceMock.checkTotalCostLimit.mockResolvedValueOnce({
      allowed: false,
      current: 20,
      reason: "Key total limit exceeded",
    });

    const session = createSession({
      key: { limitTotalUsd: 10 },
    });

    await expect(ProxyRateLimitGuard.ensure(session)).rejects.toMatchObject({
      name: "RateLimitError",
      limitType: "usd_total",
      currentUsage: 20,
      limitValue: 10,
    });
  });

  it("User 总限额超限应拦截（usd_total）", async () => {
    const { ProxyRateLimitGuard } = await import("@/app/v1/_lib/proxy/rate-limit-guard");

    rateLimitServiceMock.checkTotalCostLimit
      .mockResolvedValueOnce({ allowed: true }) // key total
      .mockResolvedValueOnce({ allowed: false, current: 20, reason: "User total limit exceeded" }); // user total

    const session = createSession({
      user: { limitTotalUsd: 10 },
    });

    await expect(ProxyRateLimitGuard.ensure(session)).rejects.toMatchObject({
      name: "RateLimitError",
      limitType: "usd_total",
      currentUsage: 20,
      limitValue: 10,
    });
  });

  it("Key 并发 Session 超限应拦截（concurrent_sessions）", async () => {
    const { ProxyRateLimitGuard } = await import("@/app/v1/_lib/proxy/rate-limit-guard");

    rateLimitServiceMock.checkAndTrackKeyUserSession.mockResolvedValueOnce({
      allowed: false,
      rejectedBy: "key",
      reasonCode: "RATE_LIMIT_CONCURRENT_SESSIONS_EXCEEDED",
      reasonParams: { current: 2, limit: 1, target: "key" },
      keyCount: 2,
      userCount: 0,
      trackedKey: false,
      trackedUser: false,
    });

    const session = createSession({
      key: { limitConcurrentSessions: 1 },
    });

    await expect(ProxyRateLimitGuard.ensure(session)).rejects.toMatchObject({
      name: "RateLimitError",
      limitType: "concurrent_sessions",
      currentUsage: 2,
      limitValue: 1,
    });
  });

  it("User 并发 Session 超限应拦截（concurrent_sessions）", async () => {
    const { ProxyRateLimitGuard } = await import("@/app/v1/_lib/proxy/rate-limit-guard");

    rateLimitServiceMock.checkAndTrackKeyUserSession.mockResolvedValueOnce({
      allowed: false,
      rejectedBy: "user",
      reasonCode: "RATE_LIMIT_CONCURRENT_SESSIONS_EXCEEDED",
      reasonParams: { current: 2, limit: 1, target: "user" },
      keyCount: 0,
      userCount: 2,
      trackedKey: false,
      trackedUser: false,
    });

    const session = createSession({
      user: { limitConcurrentSessions: 1 },
      key: { limitConcurrentSessions: 10 },
    });

    await expect(ProxyRateLimitGuard.ensure(session)).rejects.toMatchObject({
      name: "RateLimitError",
      limitType: "concurrent_sessions",
      currentUsage: 2,
      limitValue: 1,
    });
  });

  it("当 Key 并发未设置（0）且 User 并发已设置时，Key 并发检查应继承 User 并发上限", async () => {
    const { ProxyRateLimitGuard } = await import("@/app/v1/_lib/proxy/rate-limit-guard");

    const session = createSession({
      user: { limitConcurrentSessions: 15 },
      key: { limitConcurrentSessions: 0 },
    });

    await expect(ProxyRateLimitGuard.ensure(session)).resolves.toBeUndefined();

    expect(rateLimitServiceMock.checkAndTrackKeyUserSession).toHaveBeenCalledWith(
      2,
      1,
      "sess_test",
      15,
      15
    );
  });

  it("User RPM 超限应拦截（rpm）", async () => {
    const { ProxyRateLimitGuard } = await import("@/app/v1/_lib/proxy/rate-limit-guard");

    rateLimitServiceMock.checkRpmLimit.mockResolvedValueOnce({
      allowed: false,
      current: 10,
      reason: "用户每分钟请求数上限已达到（10/5）",
    });

    const session = createSession({
      user: { rpm: 5 },
    });

    await expect(ProxyRateLimitGuard.ensure(session)).rejects.toMatchObject({
      name: "RateLimitError",
      limitType: "rpm",
      currentUsage: 10,
      limitValue: 5,
    });
  });

  it("Key 5h 超限应拦截（usd_5h）", async () => {
    const { ProxyRateLimitGuard } = await import("@/app/v1/_lib/proxy/rate-limit-guard");

    rateLimitServiceMock.checkCostLimitsWithLease.mockResolvedValueOnce({
      allowed: false,
      reason: "Key 5h cost limit reached (usage: 20.0000/10.0000)",
    });

    const session = createSession({
      key: { limit5hUsd: 10 },
    });

    await expect(ProxyRateLimitGuard.ensure(session)).rejects.toMatchObject({
      name: "RateLimitError",
      limitType: "usd_5h",
      currentUsage: 20,
      limitValue: 10,
    });
  });

  it("User 5h 超限应拦截（usd_5h）", async () => {
    const { ProxyRateLimitGuard } = await import("@/app/v1/_lib/proxy/rate-limit-guard");

    rateLimitServiceMock.checkCostLimitsWithLease
      .mockResolvedValueOnce({ allowed: true }) // key 5h
      .mockResolvedValueOnce({
        allowed: false,
        reason: "User 5h cost limit reached (usage: 20.0000/10.0000)",
      }); // user 5h

    const session = createSession({
      user: { limit5hUsd: 10 },
    });

    await expect(ProxyRateLimitGuard.ensure(session)).rejects.toMatchObject({
      name: "RateLimitError",
      limitType: "usd_5h",
      currentUsage: 20,
      limitValue: 10,
    });
  });

  it("Key 周限额超限应拦截（usd_weekly）", async () => {
    const { ProxyRateLimitGuard } = await import("@/app/v1/_lib/proxy/rate-limit-guard");

    rateLimitServiceMock.checkCostLimitsWithLease
      .mockResolvedValueOnce({ allowed: true }) // key 5h
      .mockResolvedValueOnce({ allowed: true }) // user 5h
      .mockResolvedValueOnce({ allowed: true }) // key daily
      .mockResolvedValueOnce({ allowed: true }) // user daily (new with lease migration)
      .mockResolvedValueOnce({
        allowed: false,
        reason: "Key weekly cost limit reached (usage: 100.0000/10.0000)",
      }); // key weekly

    const session = createSession({
      key: { limitWeeklyUsd: 10 },
    });

    await expect(ProxyRateLimitGuard.ensure(session)).rejects.toMatchObject({
      name: "RateLimitError",
      limitType: "usd_weekly",
      currentUsage: 100,
      limitValue: 10,
    });
  });

  it("User 周限额超限应拦截（usd_weekly）", async () => {
    const { ProxyRateLimitGuard } = await import("@/app/v1/_lib/proxy/rate-limit-guard");

    rateLimitServiceMock.checkCostLimitsWithLease
      .mockResolvedValueOnce({ allowed: true }) // key 5h
      .mockResolvedValueOnce({ allowed: true }) // user 5h
      .mockResolvedValueOnce({ allowed: true }) // key daily
      .mockResolvedValueOnce({ allowed: true }) // user daily (new with lease migration)
      .mockResolvedValueOnce({ allowed: true }) // key weekly
      .mockResolvedValueOnce({
        allowed: false,
        reason: "User weekly cost limit reached (usage: 100.0000/10.0000)",
      }); // user weekly

    const session = createSession({
      user: { limitWeeklyUsd: 10 },
    });

    await expect(ProxyRateLimitGuard.ensure(session)).rejects.toMatchObject({
      name: "RateLimitError",
      limitType: "usd_weekly",
      currentUsage: 100,
      limitValue: 10,
    });
  });

  it("Key 月限额超限应拦截（usd_monthly）", async () => {
    const { ProxyRateLimitGuard } = await import("@/app/v1/_lib/proxy/rate-limit-guard");

    rateLimitServiceMock.checkCostLimitsWithLease
      .mockResolvedValueOnce({ allowed: true }) // key 5h
      .mockResolvedValueOnce({ allowed: true }) // user 5h
      .mockResolvedValueOnce({ allowed: true }) // key daily
      .mockResolvedValueOnce({ allowed: true }) // user daily (new with lease migration)
      .mockResolvedValueOnce({ allowed: true }) // key weekly
      .mockResolvedValueOnce({ allowed: true }) // user weekly
      .mockResolvedValueOnce({
        allowed: false,
        reason: "Key monthly cost limit reached (usage: 200.0000/10.0000)",
      }); // key monthly

    const session = createSession({
      key: { limitMonthlyUsd: 10 },
    });

    await expect(ProxyRateLimitGuard.ensure(session)).rejects.toMatchObject({
      name: "RateLimitError",
      limitType: "usd_monthly",
      currentUsage: 200,
      limitValue: 10,
    });
  });

  it("User 月限额超限应拦截（usd_monthly）", async () => {
    const { ProxyRateLimitGuard } = await import("@/app/v1/_lib/proxy/rate-limit-guard");

    rateLimitServiceMock.checkCostLimitsWithLease
      .mockResolvedValueOnce({ allowed: true }) // key 5h
      .mockResolvedValueOnce({ allowed: true }) // user 5h
      .mockResolvedValueOnce({ allowed: true }) // key daily
      .mockResolvedValueOnce({ allowed: true }) // user daily (new with lease migration)
      .mockResolvedValueOnce({ allowed: true }) // key weekly
      .mockResolvedValueOnce({ allowed: true }) // user weekly
      .mockResolvedValueOnce({ allowed: true }) // key monthly
      .mockResolvedValueOnce({
        allowed: false,
        reason: "User monthly cost limit reached (usage: 200.0000/10.0000)",
      }); // user monthly

    const session = createSession({
      user: { limitMonthlyUsd: 10 },
    });

    await expect(ProxyRateLimitGuard.ensure(session)).rejects.toMatchObject({
      name: "RateLimitError",
      limitType: "usd_monthly",
      currentUsage: 200,
      limitValue: 10,
    });
  });

  it("所有限额均未触发时应放行", async () => {
    const { ProxyRateLimitGuard } = await import("@/app/v1/_lib/proxy/rate-limit-guard");

    const session = createSession();
    await expect(ProxyRateLimitGuard.ensure(session)).resolves.toBeUndefined();
  });

  it("当 sessionId 缺失时，应兜底生成并继续并发检查", async () => {
    const { ProxyRateLimitGuard } = await import("@/app/v1/_lib/proxy/rate-limit-guard");

    const session = createSession() as any;
    session.sessionId = undefined;

    await expect(ProxyRateLimitGuard.ensure(session)).resolves.toBeUndefined();

    expect(generateSessionIdMock).toHaveBeenCalledTimes(1);
    expect(session.sessionId).toBe("sess_generated");
    expect(rateLimitServiceMock.checkAndTrackKeyUserSession).toHaveBeenCalledWith(
      2,
      1,
      "sess_generated",
      expect.any(Number),
      expect.any(Number)
    );
  });

  it("User daily (rolling mode) 超限应使用 checkCostLimitsWithLease", async () => {
    const { ProxyRateLimitGuard } = await import("@/app/v1/_lib/proxy/rate-limit-guard");

    rateLimitServiceMock.checkCostLimitsWithLease
      .mockResolvedValueOnce({ allowed: true }) // key 5h
      .mockResolvedValueOnce({ allowed: true }) // user 5h
      .mockResolvedValueOnce({ allowed: true }) // key daily (limit null)
      .mockResolvedValueOnce({
        allowed: false,
        reason: "User daily cost limit reached (usage: 15.0000/10.0000)",
      }); // user daily rolling

    const session = createSession({
      user: { dailyQuota: 10, dailyResetMode: "rolling", dailyResetTime: "12:00" },
      key: { limitDailyUsd: null },
    });

    await expect(ProxyRateLimitGuard.ensure(session)).rejects.toMatchObject({
      name: "RateLimitError",
      limitType: "daily_quota",
      currentUsage: 15,
      limitValue: 10,
      resetTime: null, // rolling 模式没有固定重置时间
    });

    // Verify checkCostLimitsWithLease was called with rolling mode
    expect(rateLimitServiceMock.checkCostLimitsWithLease).toHaveBeenCalledWith(1, "user", {
      limit_5h_usd: null,
      limit_daily_usd: 10,
      daily_reset_time: "12:00",
      daily_reset_mode: "rolling",
      limit_weekly_usd: null,
      limit_monthly_usd: null,
      cost_reset_at: null,
    });

    // checkUserDailyCost should NOT be called (migrated to lease)
    expect(rateLimitServiceMock.checkUserDailyCost).not.toHaveBeenCalled();
  });

  it("User daily 检查顺序：Key daily 先于 User daily", async () => {
    const { ProxyRateLimitGuard } = await import("@/app/v1/_lib/proxy/rate-limit-guard");

    const callOrder: string[] = [];

    rateLimitServiceMock.checkCostLimitsWithLease.mockImplementation(async (_id, type, limits) => {
      if (limits.limit_daily_usd !== null) {
        callOrder.push(`${type}_daily`);
      }
      return { allowed: true };
    });

    const session = createSession({
      user: { dailyQuota: 10 },
      key: { limitDailyUsd: 20 },
    });

    await ProxyRateLimitGuard.ensure(session);

    // Key daily should be checked before User daily
    const keyDailyIdx = callOrder.indexOf("key_daily");
    const userDailyIdx = callOrder.indexOf("user_daily");
    expect(keyDailyIdx).toBeLessThan(userDailyIdx);
  });
});