digikedai-bot/tests/rate-limiter.test.ts

import { describe, it, expect, vi } from "vitest";
import { RateLimiter } from "../src/core/rate-limiter.js";
import {
  MessageService,
  RATE_LIMITED_REPLY,
} from "../src/core/message-service.js";
import type { Agent } from "../src/ai/agent/agent.js";
import type { Db } from "../src/db/db.js";
import type { Logger } from "../src/utils/logger.js";

describe("RateLimiter", () => {
  it("allows up to maxPerWindow messages in a window", () => {
    const rl = new RateLimiter({
      windowMs: 60_000,
      maxPerWindow: 3,
      minIntervalMs: 0,
    });
    expect(rl.check("telegram", "u1", 0).allowed).toBe(true);
    expect(rl.check("telegram", "u1", 1).allowed).toBe(true);
    expect(rl.check("telegram", "u1", 2).allowed).toBe(true);
    expect(rl.check("telegram", "u1", 3).allowed).toBe(false);
  });

  it("enforces minIntervalMs between messages", () => {
    const rl = new RateLimiter({
      windowMs: 60_000,
      maxPerWindow: 10,
      minIntervalMs: 3_000,
    });
    expect(rl.check("telegram", "u1", 0).allowed).toBe(true);
    const d = rl.check("telegram", "u1", 1_000);
    expect(d.allowed).toBe(false);
    if (!d.allowed) expect(d.retryAfterMs).toBe(2_000);
  });

  it("isolates users by channel AND id", () => {
    const rl = new RateLimiter({
      windowMs: 60_000,
      maxPerWindow: 1,
      minIntervalMs: 0,
    });
    expect(rl.check("telegram", "u1", 0).allowed).toBe(true);
    // same id on another channel — independent budget
    expect(rl.check("shopee", "u1", 1).allowed).toBe(true);
    // different id on the same channel — independent budget
    expect(rl.check("telegram", "u2", 2).allowed).toBe(true);
  });

  it("resets after the window elapses", () => {
    const rl = new RateLimiter({
      windowMs: 60_000,
      maxPerWindow: 1,
      minIntervalMs: 0,
    });
    expect(rl.check("telegram", "u1", 0).allowed).toBe(true);
    expect(rl.check("telegram", "u1", 30_000).allowed).toBe(false);
    expect(rl.check("telegram", "u1", 60_000).allowed).toBe(true);
  });

  it("notifies at most once per window", () => {
    const rl = new RateLimiter({
      windowMs: 60_000,
      maxPerWindow: 2,
      minIntervalMs: 0,
    });
    expect(rl.check("telegram", "u1", 0).allowed).toBe(true);
    expect(rl.check("telegram", "u1", 1).allowed).toBe(true);

    let d = rl.check("telegram", "u1", 2);
    expect(d.allowed).toBe(false);
    if (!d.allowed) expect(d.notify).toBe(true);

    // further denials in the same window stay silent
    d = rl.check("telegram", "u1", 3);
    expect(d.allowed).toBe(false);
    if (!d.allowed) expect(d.notify).toBe(false);

    // new window -> quota refilled; the notification budget is restored too
    expect(rl.check("telegram", "u1", 60_000).allowed).toBe(true);
    expect(rl.check("telegram", "u1", 60_001).allowed).toBe(true);
    d = rl.check("telegram", "u1", 60_002);
    expect(d.allowed).toBe(false);
    if (!d.allowed) expect(d.notify).toBe(true);
  });
});

describe("MessageService", () => {
  function makeService(opts: { maxPerWindow: number; minIntervalMs?: number }) {
    const saveExchange = vi.fn().mockResolvedValue({ conversationId: 1 });
    const respond = vi.fn().mockResolvedValue("Hi there!");
    const logger = {
      info: vi.fn(),
      warn: vi.fn(),
      error: vi.fn(),
      debug: vi.fn(),
      trace: vi.fn(),
      fatal: vi.fn(),
      child: vi.fn().mockReturnThis(),
    } as unknown as Logger;

    const svc = new MessageService(
      { saveExchange } as unknown as Db,
      { respond } as unknown as Agent,
      new RateLimiter({
        windowMs: 60_000,
        maxPerWindow: opts.maxPerWindow,
        minIntervalMs: opts.minIntervalMs ?? 0,
      }),
      logger,
    );
    return { svc, saveExchange, respond, logger };
  }

  it("persists and replies on allowed messages", async () => {
    const { svc, saveExchange, respond } = makeService({ maxPerWindow: 5 });
    const result = await svc.handle({
      channel: "telegram",
      externalUserId: "42",
      externalConversationId: "7",
      text: "hello",
      metadata: { languageCode: "en" },
    });

    expect(result.handled).toBe(true);
    if (result.handled) {
      expect(result.reply).toBe("Hi there!");
      expect(saveExchange).toHaveBeenCalledTimes(2); // user msg + reply
      expect(respond).toHaveBeenCalledWith({
        conversationId: 1,
        userText: "hello",
        preferredLanguage: "en",
        channel: "telegram",
      });
    }
  });

  it("rate-limits without persisting or calling the LLM", async () => {
    const { svc, saveExchange, respond } = makeService({ maxPerWindow: 1 });
    await svc.handle({
      channel: "telegram",
      externalUserId: "42",
      text: "first",
    });

    saveExchange.mockClear();
    respond.mockClear();

    const result = await svc.handle({
      channel: "telegram",
      externalUserId: "42",
      text: "spam",
    });

    expect(result.handled).toBe(false);
    if (!result.handled) {
      expect(result.reason).toBe("rate_limited");
      expect(result.notify).toBe(true);
    }
    expect(saveExchange).not.toHaveBeenCalled();
    expect(respond).not.toHaveBeenCalled();
  });

  it("skips empty text without persisting or LLM", async () => {
    const { svc, saveExchange, respond } = makeService({ maxPerWindow: 5 });
    const result = await svc.handle({
      channel: "telegram",
      externalUserId: "42",
      text: "   ",
    });
    expect(result.handled).toBe(true);
    if (result.handled) expect(result.reply).toBe("");
    expect(saveExchange).not.toHaveBeenCalled();
    expect(respond).not.toHaveBeenCalled();
  });

  it("exposes a static rate-limit reply", () => {
    expect(RATE_LIMITED_REPLY.length).toBeGreaterThan(0);
  });
});