updated files

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<string, unknown> = {}
+  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)
+  })
+})

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',

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<string, ToolLimitConfig> = {
+  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<RateLimitDecision> {
+  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<void> {
+  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.
+  }
+}