sh-mcp/packages/shared/src/rate-limit.test.ts

import { describe, it, expect } from 'vitest';

import { InMemoryRateLimiter, NoopRateLimiter, RateLimitError } from './rate-limit.js';

describe('InMemoryRateLimiter', () => {
  it('allows up to the per-tool limit then throws within the window', () => {
    const limiter = new InMemoryRateLimiter({ sessionCap: 100, perToolLimit: 3, windowMs: 1000 });
    limiter.check('u', 't');
    limiter.check('u', 't');
    limiter.check('u', 't');
    expect(() => limiter.check('u', 't')).toThrow(RateLimitError);
  });

  it('isolates the per-tool window per (sub, tool)', () => {
    const limiter = new InMemoryRateLimiter({ sessionCap: 100, perToolLimit: 1, windowMs: 1000 });
    limiter.check('u', 'a');
    // Different tool — own bucket.
    expect(() => limiter.check('u', 'b')).not.toThrow();
    // Different user — own bucket.
    expect(() => limiter.check('v', 'a')).not.toThrow();
  });

  it('enforces the per-session cap across all tools', () => {
    const limiter = new InMemoryRateLimiter({ sessionCap: 2, perToolLimit: 100, windowMs: 1000 });
    limiter.check('u', 'a');
    limiter.check('u', 'b');
    expect(() => limiter.check('u', 'c')).toThrow(RateLimitError);
  });

  it('frees the per-tool window as time advances', () => {
    let t = 1_000;
    const limiter = new InMemoryRateLimiter({
      sessionCap: 100,
      perToolLimit: 1,
      windowMs: 1000,
      now: () => t,
    });
    limiter.check('u', 't');
    t += 1001; // window elapsed
    expect(() => limiter.check('u', 't')).not.toThrow();
  });

  it('carries retryAfterMs on the per-tool error', () => {
    const limiter = new InMemoryRateLimiter({ sessionCap: 100, perToolLimit: 1, windowMs: 5000 });
    limiter.check('u', 't');
    try {
      limiter.check('u', 't');
      expect.unreachable('should have thrown');
    } catch (err) {
      expect(err).toBeInstanceOf(RateLimitError);
      expect((err as RateLimitError).retryAfterMs).toBeGreaterThan(0);
    }
  });
});

describe('NoopRateLimiter', () => {
  it('never throws', () => {
    const limiter = new NoopRateLimiter();
    for (let i = 0; i < 50; i++) limiter.check('u', 't');
    expect(true).toBe(true);
  });
});