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

/**
 * In-memory rate limiting for the dispatch path.
 *
 * design.md §7.3 (Security / abuse): "per-session tool-call cap + per-tool rate
 * limit". This is the injected limiter the dispatcher consults before running a
 * handler. An in-memory implementation is sufficient for this PR (single
 * process); a distributed limiter (e.g. DynamoDB/Redis token bucket) is a
 * later-phase concern once the servers are horizontally scaled.
 *
 * Two independent limits, both keyed by the caller `sub`:
 *   1. A per-session total tool-call cap (blast-radius bound on any one user/
 *      session — design.md §2.5 prompt-injection containment).
 *   2. A per-(sub, tool) sliding-window rate limit.
 */

/** Thrown when a limit is exceeded. Adapters map this to HTTP 429. */
export class RateLimitError extends Error {
  readonly retryAfterMs: number | undefined;
  constructor(message: string, retryAfterMs?: number) {
    super(message);
    this.name = 'RateLimitError';
    this.retryAfterMs = retryAfterMs;
    Object.setPrototypeOf(this, new.target.prototype);
  }
}

/**
 * The limiter contract the dispatcher depends on. `check` throws
 * {@link RateLimitError} when the call must be rejected, otherwise records the
 * call and returns.
 */
export interface RateLimiter {
  /** Record-and-check a single tool call for `sub`/`tool`. Throws on limit. */
  check(sub: string, tool: string): void;
}

export interface RateLimitConfig {
  /** Max total tool calls a single `sub` may make for the limiter's lifetime. */
  sessionCap: number;
  /** Max calls to one tool per `sub` within {@link windowMs}. */
  perToolLimit: number;
  /** Sliding-window length in milliseconds for the per-tool limit. */
  windowMs: number;
  /** Injectable clock for deterministic tests; defaults to `Date.now`. */
  now?: () => number;
}

/**
 * Process-local {@link RateLimiter}. Construct one per server (or per test).
 *
 * Note: not shared across processes — acceptable for this PR per design.md §7.3
 * ("in-memory token bucket is fine for this PR").
 */
export class InMemoryRateLimiter implements RateLimiter {
  private readonly sessionCount = new Map<string, number>();
  private readonly toolHits = new Map<string, number[]>();
  private readonly now: () => number;

  constructor(private readonly config: RateLimitConfig) {
    this.now = config.now ?? ((): number => Date.now());
  }

  check(sub: string, tool: string): void {
    // 1. Per-session total cap.
    const sessionTotal = (this.sessionCount.get(sub) ?? 0) + 1;
    if (sessionTotal > this.config.sessionCap) {
      throw new RateLimitError(
        `Per-session tool-call cap of ${this.config.sessionCap} exceeded for "${sub}".`,
      );
    }

    // 2. Per-(sub, tool) sliding window.
    const key = `${sub}${tool}`;
    const t = this.now();
    const windowStart = t - this.config.windowMs;
    const hits = (this.toolHits.get(key) ?? []).filter((ts) => ts > windowStart);
    if (hits.length >= this.config.perToolLimit) {
      const oldest = hits[0]!;
      const retryAfterMs = oldest + this.config.windowMs - t;
      throw new RateLimitError(
        `Rate limit of ${this.config.perToolLimit} calls/${this.config.windowMs}ms to ` +
          `"${tool}" exceeded for "${sub}".`,
        retryAfterMs,
      );
    }

    // Commit both counters only after both checks pass.
    hits.push(t);
    this.toolHits.set(key, hits);
    this.sessionCount.set(sub, sessionTotal);
  }
}

/** A limiter that never throws — for tests/contexts that don't exercise limits. */
export class NoopRateLimiter implements RateLimiter {
  check(_sub: string, _tool: string): void {
    // intentionally empty
  }
}