From 37997147531ceed94547c1fb8dc849829e61687d Mon Sep 17 00:00:00 2001 From: Swayymalcolm99 Date: Sat, 17 Oct 2026 09:47:05 +0100 Subject: [PATCH] updated files Changes made --- src/__tests__/mcpRateLimit.test.ts | 214 +++++++++++++++++++++++++++++ src/mcp/codes.ts | 4 +- src/mcp/rateLimiter.ts | 141 +++++++++++++++++++ src/mcp/server.ts | 50 ++++--- src/metrics.ts | 7 + 5 files changed, 399 insertions(+), 17 deletions(-) create mode 100644 src/__tests__/mcpRateLimit.test.ts create mode 100644 src/mcp/rateLimiter.ts diff --git a/src/__tests__/mcpRateLimit.test.ts b/src/__tests__/mcpRateLimit.test.ts new file mode 100644 index 00000000..15051c31 --- /dev/null +++ b/src/__tests__/mcpRateLimit.test.ts @@ -0,0 +1,214 @@ +import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest' + +const { mockExec, mockDecr } = vi.hoisted(() => ({ + mockExec: vi.fn(), + mockDecr: vi.fn() +})) + +function chain() { + const self: Record = {} + self.incr = () => self + self.expire = () => self + self.exec = mockExec + return self +} + +vi.mock('../redis', () => ({ + redis: { + multi: () => chain(), + decr: mockDecr + }, +})) + +import { checkToolLimits, markToolCallComplete, toolRateLimit, globalConcurrencyCap } from '../mcp/rateLimiter' +import { tool_call_rejections_total } from '../metrics' + +/** ioredis multi().exec() returns [error, value] pairs in command order. */ +function counts(rateCount: number, concurrencyCount: number) { + return [ + [null, rateCount], + [null, 1], + [null, concurrencyCount], + ] +} + +beforeEach(() => { + mockExec.mockReset() + mockDecr.mockReset() + vi.clearAllMocks() +}) + +afterEach(() => { + delete process.env.MCP_TOOL_CONCURRENCY_CAP + delete process.env.MCP_PAID_CALL_RATE_LIMIT + delete process.env.MCP_BAZAAR_SEARCH_RATE_LIMIT +}) + +describe('checkToolLimits', () => { + it('allows a call inside both limits', async () => { + mockExec.mockResolvedValue(counts(1, 1)) + + expect(await checkToolLimits('paid_call')).toEqual({ allowed: true }) + }) + + it('rejects once the rate limit is exceeded', async () => { + mockExec.mockResolvedValue(counts(toolRateLimit('paid_call') + 1, 1)) + + const decision = await checkToolLimits('paid_call') + + expect(decision.allowed).toBe(false) + expect(decision.scope).toBe('rate') + expect(decision.retryAfterSeconds).toBe(60) + }) + + it('rejects once the concurrency limit is exceeded', async () => { + mockExec.mockResolvedValue(counts(1, globalConcurrencyCap() + 1)) + + const decision = await checkToolLimits('paid_call') + + expect(decision.allowed).toBe(false) + expect(decision.scope).toBe('concurrency') + expect(decision.retryAfterSeconds).toBe(1) + }) + + it('counts the attempt even when rejected by rate limit', async () => { + mockExec.mockResolvedValue(counts(toolRateLimit('paid_call') + 5, 1)) + + await checkToolLimits('paid_call') + + expect(mockExec).toHaveBeenCalledTimes(1) + }) + + it('decrements concurrency counter when rejected by concurrency limit', async () => { + mockExec.mockResolvedValue(counts(1, globalConcurrencyCap() + 1)) + + await checkToolLimits('paid_call') + + expect(mockDecr).toHaveBeenCalledTimes(1) + }) + + it('fails closed when Redis is unreachable', async () => { + mockExec.mockRejectedValue(new Error('ECONNREFUSED')) + + expect(await checkToolLimits('paid_call')).toEqual({ allowed: false, scope: 'unavailable' }) + }) + + it('fails closed when the transaction returns no result', async () => { + mockExec.mockResolvedValue(null) + + expect(await checkToolLimits('paid_call')).toEqual({ allowed: false, scope: 'unavailable' }) + }) + + it('increments rejection metric on rate limit', async () => { + mockExec.mockResolvedValue(counts(toolRateLimit('paid_call') + 1, 1)) + + await checkToolLimits('paid_call') + + expect(tool_call_rejections_total.inc).toHaveBeenCalledWith({ tool: 'paid_call', reason: 'rate_limit' }) + }) + + it('increments rejection metric on concurrency limit', async () => { + mockExec.mockResolvedValue(counts(1, globalConcurrencyCap() + 1)) + + await checkToolLimits('paid_call') + + expect(tool_call_rejections_total.inc).toHaveBeenCalledWith({ tool: 'paid_call', reason: 'concurrency_limit' }) + }) + + it('increments rejection metric on Redis unavailability', async () => { + mockExec.mockRejectedValue(new Error('ECONNREFUSED')) + + await checkToolLimits('paid_call') + + expect(tool_call_rejections_total.inc).toHaveBeenCalledWith({ tool: 'paid_call', reason: 'unavailable' }) + }) +}) + +describe('toolRateLimit', () => { + it('returns default limit for paid_call', () => { + expect(toolRateLimit('paid_call')).toBe(10) + }) + + it('returns default limit for bazaar_search', () => { + expect(toolRateLimit('bazaar_search')).toBe(20) + }) + + it('returns global default for unknown tools', () => { + expect(toolRateLimit('unknown_tool')).toBe(30) + }) + + it('reads limit from environment variable', () => { + process.env.MCP_PAID_CALL_RATE_LIMIT = '5' + + expect(toolRateLimit('paid_call')).toBe(5) + }) + + it('falls back to default for invalid env value', () => { + process.env.MCP_PAID_CALL_RATE_LIMIT = 'invalid' + + expect(toolRateLimit('paid_call')).toBe(10) + }) +}) + +describe('globalConcurrencyCap', () => { + it('returns default concurrency cap', () => { + expect(globalConcurrencyCap()).toBe(10) + }) + + it('reads cap from environment variable', () => { + process.env.MCP_TOOL_CONCURRENCY_CAP = '20' + + expect(globalConcurrencyCap()).toBe(20) + }) + + it('falls back to default for invalid env value', () => { + process.env.MCP_TOOL_CONCURRENCY_CAP = 'invalid' + + expect(globalConcurrencyCap()).toBe(10) + }) +}) + +describe('markToolCallComplete', () => { + it('decrements the in-flight counter', async () => { + mockDecr.mockResolvedValue(1) + + await markToolCallComplete() + + expect(mockDecr).toHaveBeenCalledTimes(1) + expect(mockDecr).toHaveBeenCalledWith('lens:mcp:concurrency:in-flight') + }) + + it('does not throw when Redis is unreachable', async () => { + mockDecr.mockRejectedValue(new Error('ECONNREFUSED')) + + await expect(markToolCallComplete()).resolves.not.toThrow() + }) +}) + +describe('burst throttling behavior', () => { + it('throttles a burst rather than queuing indefinitely', async () => { + // Simulate a burst of 15 calls when rate limit is 10 + const results = [] + + for (let i = 0; i < 15; i++) { + const count = i + 1 + mockExec.mockResolvedValueOnce(counts(count, 1)) + results.push(await checkToolLimits('paid_call')) + } + + // First 10 should be allowed + for (let i = 0; i < 10; i++) { + expect(results[i].allowed).toBe(true) + } + + // Next 5 should be rejected with rate limit error + for (let i = 10; i < 15; i++) { + expect(results[i].allowed).toBe(false) + expect(results[i].scope).toBe('rate') + expect(results[i].retryAfterSeconds).toBe(60) + } + + // Verify rejections were counted + expect(tool_call_rejections_total.inc).toHaveBeenCalledTimes(5) + }) +}) diff --git a/src/mcp/codes.ts b/src/mcp/codes.ts index a9631225..973ed96e 100644 --- a/src/mcp/codes.ts +++ b/src/mcp/codes.ts @@ -6,7 +6,9 @@ export const MCP_ERROR_CODES = { ERR_SERVICE_UNREACHABLE: 'ERR_SERVICE_UNREACHABLE', ERR_NO_RESULTS: 'ERR_NO_RESULTS', ERR_BAD_REQUEST: 'ERR_BAD_REQUEST', - ERR_INTERNAL_FAILURE: 'ERR_INTERNAL_FAILURE' + ERR_INTERNAL_FAILURE: 'ERR_INTERNAL_FAILURE', + ERR_RATE_LIMIT_EXCEEDED: 'ERR_RATE_LIMIT_EXCEEDED', + ERR_CONCURRENCY_LIMIT_EXCEEDED: 'ERR_CONCURRENCY_LIMIT_EXCEEDED' } as const; export type McpErrorCode = keyof typeof MCP_ERROR_CODES; diff --git a/src/mcp/rateLimiter.ts b/src/mcp/rateLimiter.ts new file mode 100644 index 00000000..a3638546 --- /dev/null +++ b/src/mcp/rateLimiter.ts @@ -0,0 +1,141 @@ +import type Redis from 'ioredis' +import { redis } from '../redis' +import { tool_call_rejections_total } from '../metrics' + +/** + * Per-tool rate limiting and concurrency control for MCP tool calls. + * + * Prevents an agent in a loop from issuing unbounded calls that cost money + * (paid_call) or create unbounded outbound connections (bazaar_search). + * + * Limits are: + * - Per-tool rate limit: configurable requests per minute per tool + * - Global concurrency cap: maximum in-flight tool calls across all tools + * + * Exceeding a limit returns a structured error with a retry hint rather than + * silently queueing, so the agent knows to slow down. + */ + +interface ToolLimitConfig { + /** Maximum requests per minute for this tool */ + requestsPerMinute: number + /** Maximum concurrent in-flight calls for this tool (optional, uses global cap if not set) */ + maxConcurrent?: number +} + +/** Default per-tool limits. Expensive tools get tighter limits. */ +const DEFAULT_TOOL_LIMITS: Record = { + paid_call: { requestsPerMinute: 10, maxConcurrent: 3 }, + bazaar_search: { requestsPerMinute: 20, maxConcurrent: 5 }, +} + +/** Global concurrency cap for all tool calls (fallback for tools without specific limits) */ +const GLOBAL_CONCURRENCY_CAP = 10 + +/** Global rate limit fallback for tools without specific limits */ +const GLOBAL_RATE_LIMIT = 30 + +function envLimit(name: string, fallback: number): number { + const raw = Number(process.env[name]) + return Number.isFinite(raw) && raw > 0 ? Math.floor(raw) : fallback +} + +/** Global concurrency cap from env or default */ +export function globalConcurrencyCap(): number { + return envLimit('MCP_TOOL_CONCURRENCY_CAP', GLOBAL_CONCURRENCY_CAP) +} + +/** Per-tool rate limit from env or default */ +export function toolRateLimit(toolName: string): number { + const envKey = `MCP_${toolName.toUpperCase()}_RATE_LIMIT` + const toolDefault = DEFAULT_TOOL_LIMITS[toolName]?.requestsPerMinute + return envLimit(envKey, toolDefault ?? GLOBAL_RATE_LIMIT) +} + +/** Per-tool concurrency limit from env or default */ +export function toolConcurrencyLimit(toolName: string): number { + const envKey = `MCP_${toolName.toUpperCase()}_CONCURRENCY_LIMIT` + const toolDefault = DEFAULT_TOOL_LIMITS[toolName]?.maxConcurrent + return envLimit(envKey, toolDefault ?? globalConcurrencyCap()) +} + +export interface RateLimitDecision { + allowed: boolean + /** Which limit rejected the call */ + scope?: 'rate' | 'concurrency' | 'unavailable' + retryAfterSeconds?: number +} + +const WINDOW_SECONDS = 60 + +function rateLimitKey(toolName: string): string { + return `lens:mcp:rate:${toolName}:${Math.floor(Date.now() / 1000 / WINDOW_SECONDS)}` +} + +function concurrencyKey(): string { + return 'lens:mcp:concurrency:in-flight' +} + +/** + * Checks if a tool call is allowed under both rate and concurrency limits. + * + * Rate limiting is per-tool and counts the attempt even when rejected, so a + * flood still costs its slots. + * + * Concurrency limiting is global across all tools to prevent total resource + * exhaustion. + * + * Fails closed: if Redis is unreachable, we reject the call to be safe. + */ +export async function checkToolLimits(toolName: string): Promise { + const db = redis as unknown as Redis + + try { + const results = await db + .multi() + .incr(rateLimitKey(toolName)) + .expire(rateLimitKey(toolName), WINDOW_SECONDS) + .incr(concurrencyKey()) + .exec() + + if (!results) return { allowed: false, scope: 'unavailable' } + + const rateCount = Number(results[0]?.[1] ?? 0) + const concurrencyCount = Number(results[2]?.[1] ?? 0) + + const rateLimit = toolRateLimit(toolName) + const concurrencyLimit = globalConcurrencyCap() + + if (rateCount > rateLimit) { + tool_call_rejections_total.inc({ tool: toolName, reason: 'rate_limit' }) + return { allowed: false, scope: 'rate', retryAfterSeconds: WINDOW_SECONDS } + } + + if (concurrencyCount > concurrencyLimit) { + tool_call_rejections_total.inc({ tool: toolName, reason: 'concurrency_limit' }) + await db.decr(concurrencyKey()) // Rollback the increment since we're rejecting + return { allowed: false, scope: 'concurrency', retryAfterSeconds: 1 } + } + + return { allowed: true } + } catch (error) { + tool_call_rejections_total.inc({ tool: toolName, reason: 'unavailable' }) + return { allowed: false, scope: 'unavailable' } + } +} + +/** + * Marks a tool call as complete, decrementing the in-flight counter. + * + * Must be called after a tool call finishes (success or failure) to free up + * the concurrency slot. + */ +export async function markToolCallComplete(): Promise { + const db = redis as unknown as Redis + try { + await db.decr(concurrencyKey()) + } catch { + // If Redis is down, we can't decrement. This is acceptable because the + // key will naturally expire and we fail closed on new calls anyway. + } +} diff --git a/src/mcp/server.ts b/src/mcp/server.ts index 3f69c5b1..55857112 100644 --- a/src/mcp/server.ts +++ b/src/mcp/server.ts @@ -3,6 +3,7 @@ import { StdioServerTransport } from '@modelcontextprotocol/sdk/server/stdio.js' import { CallToolRequestSchema, ListToolsRequestSchema } from '@modelcontextprotocol/sdk/types.js'; import { MCP_ERROR_CODES, McpErrorCode } from './codes.js'; import { fetchWithTimeout } from '../fetch.js'; +import { checkToolLimits, markToolCallComplete } from './rateLimiter.js'; export interface McpServerConfig { fetchWithPayment: typeof fetch; @@ -72,11 +73,31 @@ export class BazaarMcpServer { })); this.server.setRequestHandler(CallToolRequestSchema, async (request) => { - if (request.params.name === 'bazaar_search') { - const url = String(request.params.arguments?.url); - const query = request.params.arguments?.query ? String(request.params.arguments?.query) : undefined; + const toolName = request.params.name; + + // Check rate and concurrency limits before executing + const limitCheck = await checkToolLimits(toolName); + if (!limitCheck.allowed) { + const errorCode = limitCheck.scope === 'rate' + ? 'ERR_RATE_LIMIT_EXCEEDED' + : limitCheck.scope === 'concurrency' + ? 'ERR_CONCURRENCY_LIMIT_EXCEEDED' + : 'ERR_SERVICE_UNREACHABLE'; - try { + const reason = limitCheck.scope === 'rate' + ? `Rate limit exceeded. Retry after ${limitCheck.retryAfterSeconds}s.` + : limitCheck.scope === 'concurrency' + ? `Too many concurrent tool calls. Retry after ${limitCheck.retryAfterSeconds}s.` + : 'Rate limiting service unavailable.'; + + return this.createErrorResponse(errorCode, reason); + } + + try { + if (toolName === 'bazaar_search') { + const url = String(request.params.arguments?.url); + const query = request.params.arguments?.query ? String(request.params.arguments?.query) : undefined; + const endpoint = query ? `${url.replace(/\/$/, '')}/discovery/search?q=${encodeURIComponent(query)}` : `${url.replace(/\/$/, '')}/discovery/resources`; @@ -99,16 +120,12 @@ export class BazaarMcpServer { return { content: [{ type: 'text', text: JSON.stringify(data, null, 2) }], }; - } catch (e: any) { - return this.createErrorResponse('ERR_SERVICE_UNREACHABLE', e.message); } - } - if (request.params.name === 'paid_call') { - const url = String(request.params.arguments?.url); - const method = request.params.arguments?.method ? String(request.params.arguments?.method) : 'GET'; - - try { + if (toolName === 'paid_call') { + const url = String(request.params.arguments?.url); + const method = request.params.arguments?.method ? String(request.params.arguments?.method) : 'GET'; + const res = await this.fetchWithPayment(url, { method }); if (!res.ok) { @@ -129,12 +146,13 @@ export class BazaarMcpServer { return { content: [{ type: 'text', text: typeof parsed === 'string' ? parsed : JSON.stringify(parsed, null, 2) }], }; - } catch (e: any) { - return this.handlePaidCallError(e); } - } - throw new Error(`Unknown tool: ${request.params.name}`); + throw new Error(`Unknown tool: ${toolName}`); + } finally { + // Always mark the tool call as complete to free the concurrency slot + await markToolCallComplete(); + } }); } diff --git a/src/metrics.ts b/src/metrics.ts index 39a3a505..eb78133e 100644 --- a/src/metrics.ts +++ b/src/metrics.ts @@ -98,6 +98,13 @@ export const db_query_duration_seconds = new Histogram({ registers: [register] }) +export const tool_call_rejections_total = new Counter({ + name: 'tool_call_rejections_total', + help: 'Total number of MCP tool calls rejected due to rate or concurrency limits', + labelNames: ['tool', 'reason'], + registers: [register] +}) + export async function getMetrics() { return await register.metrics() }