How to rate-limit MCP tool calls

1// src/rate-limiter.ts
2export interface RateLimitConfig {
3  maxTokens: number;       // Maximum bucket capacity
4  refillRate: number;      // Tokens added per second
5}
6
7interface Bucket {
8  tokens: number;
9  lastRefill: number;
10}
11
12export class TokenBucketLimiter {
13  private buckets = new Map<string, Bucket>();
14  private configs = new Map<string, RateLimitConfig>();
15  private defaultConfig: RateLimitConfig;
16
17  constructor(defaultConfig: RateLimitConfig = { maxTokens: 60, refillRate: 1 }) {
18    this.defaultConfig = defaultConfig;
19  }
20
21  setConfig(key: string, config: RateLimitConfig): void {
22    this.configs.set(key, config);
23  }
24
25  tryConsume(key: string): { allowed: boolean; retryAfterMs: number; remaining: number } {
26    const config = this.configs.get(key) || this.defaultConfig;
27    const now = Date.now();
28    let bucket = this.buckets.get(key);
29
30    if (!bucket) {
31      bucket = { tokens: config.maxTokens, lastRefill: now };
32      this.buckets.set(key, bucket);
33    }
34
35    // Refill tokens based on elapsed time
36    const elapsed = (now - bucket.lastRefill) / 1000;
37    bucket.tokens = Math.min(config.maxTokens, bucket.tokens + elapsed * config.refillRate);
38    bucket.lastRefill = now;
39
40    if (bucket.tokens >= 1) {
41      bucket.tokens -= 1;
42      return { allowed: true, retryAfterMs: 0, remaining: Math.floor(bucket.tokens) };
43    }
44
45    // Calculate when next token will be available
46    const deficit = 1 - bucket.tokens;
47    const retryAfterMs = Math.ceil((deficit / config.refillRate) * 1000);
48    return { allowed: false, retryAfterMs, remaining: 0 };
49  }
50
51  reset(key: string): void {
52    this.buckets.delete(key);
53  }
54}

1// src/rate-config.ts
2import { RateLimitConfig } from "./rate-limiter.js";
3
4export const TOOL_RATE_LIMITS: Record<string, RateLimitConfig> = {
5  // Local file operations — fast and cheap
6  list_files:    { maxTokens: 120, refillRate: 2 },   // 120/min burst, 2/sec sustained
7  read_file:     { maxTokens: 60, refillRate: 1 },    // 60/min burst, 1/sec sustained
8  
9  // Database queries — moderate cost
10  query_db:      { maxTokens: 30, refillRate: 0.5 },  // 30/min burst, 30/min sustained
11  
12  // External API calls — expensive and upstream-limited
13  call_api:      { maxTokens: 10, refillRate: 0.17 }, // 10/min burst, ~10/min sustained
14  
15  // Write operations — highest risk
16  write_file:    { maxTokens: 20, refillRate: 0.33 }, // 20/min burst, ~20/min sustained
17  delete_file:   { maxTokens: 5, refillRate: 0.08 },  // 5/min burst, ~5/min sustained
18};

1// src/rate-limited-tool.ts
2import { TokenBucketLimiter } from "./rate-limiter.js";
3import { TOOL_RATE_LIMITS } from "./rate-config.js";
4
5const limiter = new TokenBucketLimiter();
6
7// Initialize tool-specific configs
8for (const [tool, config] of Object.entries(TOOL_RATE_LIMITS)) {
9  limiter.setConfig(tool, config);
10}
11
12export function withRateLimit<T>(
13  toolName: string,
14  handler: (params: T) => Promise<any>,
15  userId?: string
16): (params: T) => Promise<any> {
17  return async (params: T) => {
18    // Check per-tool limit
19    const toolResult = limiter.tryConsume(toolName);
20    if (!toolResult.allowed) {
21      console.error(`[rate-limit] Tool ${toolName} rate limited. Retry after ${toolResult.retryAfterMs}ms`);
22      return {
23        content: [{
24          type: "text",
25          text: `Rate limit exceeded for ${toolName}. Please retry after ${Math.ceil(toolResult.retryAfterMs / 1000)} seconds. Remaining capacity: ${toolResult.remaining}.`,
26        }],
27        isError: true,
28      };
29    }
30
31    // Check per-user limit if userId provided
32    if (userId) {
33      const userKey = `user:${userId}`;
34      const userResult = limiter.tryConsume(userKey);
35      if (!userResult.allowed) {
36        console.error(`[rate-limit] User ${userId} rate limited`);
37        return {
38          content: [{
39            type: "text",
40            text: `User rate limit exceeded. Please retry after ${Math.ceil(userResult.retryAfterMs / 1000)} seconds.`,
41          }],
42          isError: true,
43        };
44      }
45    }
46
47    return handler(params);
48  };
49}

1// src/index.ts
2import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
3import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js";
4import { z } from "zod";
5import { withRateLimit } from "./rate-limited-tool.js";
6
7const server = new McpServer({ name: "rate-limited-server", version: "1.0.0" });
8
9// Register tools with rate limiting
10server.tool(
11  "read_file",
12  "Read a file with rate limiting applied",
13  { filePath: z.string() },
14  withRateLimit("read_file", async ({ filePath }) => {
15    const fs = await import("fs/promises");
16    const content = await fs.readFile(filePath, "utf-8");
17    return { content: [{ type: "text", text: content }] };
18  })
19);
20
21server.tool(
22  "call_api",
23  "Call an external API with strict rate limiting",
24  { url: z.string(), method: z.enum(["GET", "POST"]).default("GET") },
25  withRateLimit("call_api", async ({ url, method }) => {
26    const response = await fetch(url, { method });
27    const text = await response.text();
28    return { content: [{ type: "text", text }] };
29  })
30);
31
32async function main() {
33  const transport = new StdioServerTransport();
34  await server.connect(transport);
35  console.error("Rate-limited MCP server running");
36}
37main().catch(e => { console.error(e); process.exit(1); });