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 4f9f185f..cb21beb2 100644 --- a/src/mcp/codes.ts +++ b/src/mcp/codes.ts @@ -6,6 +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_RATE_LIMIT_EXCEEDED: 'ERR_RATE_LIMIT_EXCEEDED', + ERR_CONCURRENCY_LIMIT_EXCEEDED: 'ERR_CONCURRENCY_LIMIT_EXCEEDED' ERR_INVALID_ASSET: 'ERR_INVALID_ASSET', ERR_PAIR_NOT_WATCHED: 'ERR_PAIR_NOT_WATCHED', ERR_ASSET_ISSUER_MISMATCH: 'ERR_ASSET_ISSUER_MISMATCH', 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 4e3ead34..04046cda 100644 --- a/src/mcp/server.ts +++ b/src/mcp/server.ts @@ -37,6 +37,7 @@ import { } from './toolSchemas.js'; import { boundSerialized, readBodyCapped, type McpLimits, type ResultBudget } from './limits.js'; import { fetchWithTimeout } from '../fetch.js'; +import { checkToolLimits, markToolCallComplete } from './rateLimiter.js'; import { config } from '../config.js'; import { ssrfSafeFetch, SsrfError, type SsrfFetchOptions } from '../net/ssrf.js'; import { defaultResources, snapshotResource, type LensResourceDescriptor } from './resources.js'; @@ -401,6 +402,38 @@ export class BazaarMcpServer { }), })); + this.server.setRequestHandler(CallToolRequestSchema, async (request) => { + 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'; + + 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`; + + const res = await ssrfSafeFetch(endpoint, {}, this.bazaarFetchOptions); + target.setRequestHandler(ReadResourceRequestSchema, async (request) => { const uri = request.params.uri; const descriptor = this.resources.find((d) => d.uri === uri); @@ -578,6 +611,10 @@ export class BazaarMcpServer { ); } + if (toolName === 'paid_call') { + const url = String(request.params.arguments?.url); + const method = request.params.arguments?.method ? String(request.params.arguments?.method) : 'GET'; + if (request.params.name === 'paid_call') { const parsed = PAID_CALL_ARGUMENTS_SCHEMA.safeParse(request.params.arguments ?? {}); if (!parsed.success) { @@ -665,6 +702,10 @@ export class BazaarMcpServer { // answer. parsed = body.text; } + + return { + content: [{ type: 'text', text: typeof parsed === 'string' ? parsed : JSON.stringify(parsed, null, 2) }], + }; return this.createBoundedResult(parsed, budget); } catch (e: any) { @@ -682,8 +723,12 @@ export class BazaarMcpServer { }); return this.handlePaidCallError(e); } - } + throw new Error(`Unknown tool: ${toolName}`); + } finally { + // Always mark the tool call as complete to free the concurrency slot + await markToolCallComplete(); + } const rawName = String(request.params.name ?? ''); const boundedName = rawName.slice(0, 64).replace(/[\r\n\t]/g, ' '); return this.createErrorResponse( diff --git a/src/metrics.ts b/src/metrics.ts index 8550de3f..b9b0ca96 100644 --- a/src/metrics.ts +++ b/src/metrics.ts @@ -244,6 +244,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] +}) + // Database connection pool utilisation metrics export const db_pool_connections_in_use = new Gauge({ name: 'db_pool_connections_in_use',