@sparkleideas/mcp 3.0.0-alpha.9

package/src/prompt-registry.ts

@@ -0,0 +1,296 @@
+/**
+ * @sparkleideas/mcp - Prompt Registry
+ *
+ * MCP 2025-11-25 compliant prompt management
+ * Supports: list, get, arguments, templates, embedded resources
+ */
+
+import { EventEmitter } from 'events';
+import type {
+  MCPPrompt,
+  PromptArgument,
+  PromptMessage,
+  PromptListResult,
+  PromptGetResult,
+  PromptContent,
+  TextContent,
+  EmbeddedResource,
+  ResourceContent,
+  ILogger,
+} from './types.js';
+
+export type PromptHandler = (
+  args: Record<string, string>
+) => Promise<PromptMessage[]>;
+
+export interface PromptDefinition extends MCPPrompt {
+  handler: PromptHandler;
+}
+
+export interface PromptRegistryOptions {
+  maxPrompts?: number;
+  validateArguments?: boolean;
+}
+
+export class PromptRegistry extends EventEmitter {
+  private prompts: Map<string, PromptDefinition> = new Map();
+
+  private readonly options: Required<PromptRegistryOptions>;
+
+  constructor(
+    private readonly logger: ILogger,
+    options: PromptRegistryOptions = {}
+  ) {
+    super();
+    this.options = {
+      maxPrompts: options.maxPrompts ?? 1000,
+      validateArguments: options.validateArguments ?? true,
+    };
+  }
+
+  /**
+   * Register a prompt
+   */
+  register(prompt: PromptDefinition): boolean {
+    if (this.prompts.size >= this.options.maxPrompts) {
+      this.logger.error('Maximum prompts reached', { max: this.options.maxPrompts });
+      return false;
+    }
+
+    if (this.prompts.has(prompt.name)) {
+      this.logger.warn('Prompt already registered', { name: prompt.name });
+      return false;
+    }
+
+    this.prompts.set(prompt.name, prompt);
+
+    this.logger.debug('Prompt registered', { name: prompt.name });
+    this.emit('prompt:registered', { name: prompt.name });
+    this.emitListChanged();
+
+    return true;
+  }
+
+  /**
+   * Unregister a prompt
+   */
+  unregister(name: string): boolean {
+    if (!this.prompts.has(name)) {
+      return false;
+    }
+
+    this.prompts.delete(name);
+
+    this.logger.debug('Prompt unregistered', { name });
+    this.emit('prompt:unregistered', { name });
+    this.emitListChanged();
+
+    return true;
+  }
+
+  /**
+   * List prompts with pagination
+   */
+  list(cursor?: string, pageSize: number = 50): PromptListResult {
+    const allPrompts = Array.from(this.prompts.values()).map((p) => ({
+      name: p.name,
+      title: p.title,
+      description: p.description,
+      arguments: p.arguments,
+    }));
+
+    let startIndex = 0;
+    if (cursor) {
+      const decoded = this.decodeCursor(cursor);
+      startIndex = decoded.offset;
+    }
+
+    const endIndex = Math.min(startIndex + pageSize, allPrompts.length);
+    const prompts = allPrompts.slice(startIndex, endIndex);
+
+    const result: PromptListResult = { prompts };
+
+    if (endIndex < allPrompts.length) {
+      result.nextCursor = this.encodeCursor({ offset: endIndex });
+    }
+
+    return result;
+  }
+
+  /**
+   * Get a prompt with arguments
+   */
+  async get(
+    name: string,
+    args: Record<string, string> = {}
+  ): Promise<PromptGetResult> {
+    const prompt = this.prompts.get(name);
+    if (!prompt) {
+      throw new Error(`Prompt not found: ${name}`);
+    }
+
+    // Validate required arguments
+    if (this.options.validateArguments && prompt.arguments) {
+      for (const arg of prompt.arguments) {
+        if (arg.required && !(arg.name in args)) {
+          throw new Error(`Missing required argument: ${arg.name}`);
+        }
+      }
+    }
+
+    const messages = await prompt.handler(args);
+
+    this.emit('prompt:get', { name, argCount: Object.keys(args).length });
+
+    return {
+      description: prompt.description,
+      messages,
+    };
+  }
+
+  /**
+   * Get prompt by name
+   */
+  getPrompt(name: string): MCPPrompt | undefined {
+    const prompt = this.prompts.get(name);
+    if (!prompt) return undefined;
+
+    return {
+      name: prompt.name,
+      title: prompt.title,
+      description: prompt.description,
+      arguments: prompt.arguments,
+    };
+  }
+
+  /**
+   * Check if prompt exists
+   */
+  hasPrompt(name: string): boolean {
+    return this.prompts.has(name);
+  }
+
+  /**
+   * Get prompt count
+   */
+  getPromptCount(): number {
+    return this.prompts.size;
+  }
+
+  /**
+   * Get stats
+   */
+  getStats(): {
+    totalPrompts: number;
+    promptsWithArgs: number;
+  } {
+    let promptsWithArgs = 0;
+    for (const prompt of this.prompts.values()) {
+      if (prompt.arguments && prompt.arguments.length > 0) {
+        promptsWithArgs++;
+      }
+    }
+
+    return {
+      totalPrompts: this.prompts.size,
+      promptsWithArgs,
+    };
+  }
+
+  /**
+   * Encode cursor for pagination
+   */
+  private encodeCursor(data: { offset: number }): string {
+    return Buffer.from(JSON.stringify(data)).toString('base64');
+  }
+
+  /**
+   * Decode cursor for pagination
+   */
+  private decodeCursor(cursor: string): { offset: number } {
+    try {
+      return JSON.parse(Buffer.from(cursor, 'base64').toString('utf-8'));
+    } catch {
+      return { offset: 0 };
+    }
+  }
+
+  /**
+   * Emit listChanged notification
+   */
+  private emitListChanged(): void {
+    this.emit('prompts:listChanged');
+  }
+}
+
+export function createPromptRegistry(
+  logger: ILogger,
+  options?: PromptRegistryOptions
+): PromptRegistry {
+  return new PromptRegistry(logger, options);
+}
+
+/**
+ * Helper to define a prompt
+ */
+export function definePrompt(
+  name: string,
+  description: string,
+  handler: PromptHandler,
+  options?: {
+    title?: string;
+    arguments?: PromptArgument[];
+  }
+): PromptDefinition {
+  return {
+    name,
+    description,
+    title: options?.title,
+    arguments: options?.arguments,
+    handler,
+  };
+}
+
+/**
+ * Helper to create a text message
+ */
+export function textMessage(
+  role: 'user' | 'assistant',
+  text: string
+): PromptMessage {
+  return {
+    role,
+    content: {
+      type: 'text',
+      text,
+    } as TextContent,
+  };
+}
+
+/**
+ * Helper to create a message with embedded resource
+ */
+export function resourceMessage(
+  role: 'user' | 'assistant',
+  resource: ResourceContent
+): PromptMessage {
+  return {
+    role,
+    content: {
+      type: 'resource',
+      resource,
+    } as EmbeddedResource,
+  };
+}
+
+/**
+ * Template string interpolation for prompts
+ */
+export function interpolate(
+  template: string,
+  args: Record<string, string>
+): string {
+  return template.replace(/\{(\w+)\}/g, (match, key) => {
+    return args[key] ?? match;
+  });
+}

package/src/rate-limiter.ts

@@ -0,0 +1,266 @@
+/**
+ * @sparkleideas/mcp - Rate Limiter
+ *
+ * Token bucket rate limiting for DoS protection
+ */
+
+import { EventEmitter } from 'events';
+import type { ILogger } from './types.js';
+
+export interface RateLimitConfig {
+  /** Requests per second */
+  requestsPerSecond: number;
+  /** Burst size (max tokens) */
+  burstSize: number;
+  /** Per-session limits (if different from global) */
+  perSessionLimit?: number;
+  /** Cleanup interval for expired sessions */
+  cleanupInterval?: number;
+}
+
+export interface RateLimitResult {
+  allowed: boolean;
+  remaining: number;
+  resetIn: number; // ms until bucket refills
+  retryAfter?: number; // seconds to wait (for 429 response)
+}
+
+interface TokenBucket {
+  tokens: number;
+  lastRefill: number;
+}
+
+const DEFAULT_CONFIG: RateLimitConfig = {
+  requestsPerSecond: 100,
+  burstSize: 200,
+  perSessionLimit: 50,
+  cleanupInterval: 60000, // 1 minute
+};
+
+export class RateLimiter extends EventEmitter {
+  private readonly config: Required<RateLimitConfig>;
+  private globalBucket: TokenBucket;
+  private sessionBuckets: Map<string, TokenBucket> = new Map();
+  private cleanupTimer?: NodeJS.Timeout;
+
+  constructor(
+    private readonly logger: ILogger,
+    config: Partial<RateLimitConfig> = {}
+  ) {
+    super();
+    this.config = { ...DEFAULT_CONFIG, ...config } as Required<RateLimitConfig>;
+
+    // Initialize global bucket
+    this.globalBucket = {
+      tokens: this.config.burstSize,
+      lastRefill: Date.now(),
+    };
+
+    this.startCleanup();
+  }
+
+  /**
+   * Check if request is allowed (global limit)
+   */
+  checkGlobal(): RateLimitResult {
+    return this.checkBucket(this.globalBucket, this.config.requestsPerSecond, this.config.burstSize);
+  }
+
+  /**
+   * Check if request is allowed (per-session limit)
+   */
+  checkSession(sessionId: string): RateLimitResult {
+    let bucket = this.sessionBuckets.get(sessionId);
+    if (!bucket) {
+      bucket = {
+        tokens: this.config.perSessionLimit,
+        lastRefill: Date.now(),
+      };
+      this.sessionBuckets.set(sessionId, bucket);
+    }
+
+    return this.checkBucket(
+      bucket,
+      this.config.perSessionLimit / 10, // Refill rate (10 seconds to full)
+      this.config.perSessionLimit
+    );
+  }
+
+  /**
+   * Check both global and session limits
+   */
+  check(sessionId?: string): RateLimitResult {
+    const globalResult = this.checkGlobal();
+    if (!globalResult.allowed) {
+      this.emit('rate-limit:global', { remaining: globalResult.remaining });
+      return globalResult;
+    }
+
+    if (sessionId) {
+      const sessionResult = this.checkSession(sessionId);
+      if (!sessionResult.allowed) {
+        this.emit('rate-limit:session', { sessionId, remaining: sessionResult.remaining });
+        return sessionResult;
+      }
+      return sessionResult;
+    }
+
+    return globalResult;
+  }
+
+  /**
+   * Consume a token (call after request is processed)
+   */
+  consume(sessionId?: string): void {
+    this.consumeFromBucket(this.globalBucket);
+    if (sessionId) {
+      const bucket = this.sessionBuckets.get(sessionId);
+      if (bucket) {
+        this.consumeFromBucket(bucket);
+      }
+    }
+  }
+
+  /**
+   * Reset session bucket
+   */
+  resetSession(sessionId: string): void {
+    this.sessionBuckets.delete(sessionId);
+    this.logger.debug('Rate limit session reset', { sessionId });
+  }
+
+  /**
+   * Get current stats
+   */
+  getStats(): {
+    globalTokens: number;
+    globalBurstSize: number;
+    sessionCount: number;
+    config: RateLimitConfig;
+  } {
+    this.refillBucket(this.globalBucket, this.config.requestsPerSecond, this.config.burstSize);
+    return {
+      globalTokens: Math.floor(this.globalBucket.tokens),
+      globalBurstSize: this.config.burstSize,
+      sessionCount: this.sessionBuckets.size,
+      config: this.config,
+    };
+  }
+
+  /**
+   * Destroy the rate limiter
+   */
+  destroy(): void {
+    if (this.cleanupTimer) {
+      clearInterval(this.cleanupTimer);
+      this.cleanupTimer = undefined;
+    }
+    this.sessionBuckets.clear();
+    this.removeAllListeners();
+  }
+
+  /**
+   * Check bucket and refill tokens
+   */
+  private checkBucket(
+    bucket: TokenBucket,
+    refillRate: number,
+    maxTokens: number
+  ): RateLimitResult {
+    this.refillBucket(bucket, refillRate, maxTokens);
+
+    if (bucket.tokens >= 1) {
+      return {
+        allowed: true,
+        remaining: Math.floor(bucket.tokens) - 1,
+        resetIn: Math.ceil((maxTokens - bucket.tokens) / refillRate * 1000),
+      };
+    }
+
+    // Calculate when bucket will have 1 token
+    const tokensNeeded = 1 - bucket.tokens;
+    const retryAfter = Math.ceil(tokensNeeded / refillRate);
+
+    return {
+      allowed: false,
+      remaining: 0,
+      resetIn: retryAfter * 1000,
+      retryAfter,
+    };
+  }
+
+  /**
+   * Refill tokens based on elapsed time
+   */
+  private refillBucket(bucket: TokenBucket, refillRate: number, maxTokens: number): void {
+    const now = Date.now();
+    const elapsed = (now - bucket.lastRefill) / 1000; // seconds
+    const tokensToAdd = elapsed * refillRate;
+
+    bucket.tokens = Math.min(maxTokens, bucket.tokens + tokensToAdd);
+    bucket.lastRefill = now;
+  }
+
+  /**
+   * Consume a token from bucket
+   */
+  private consumeFromBucket(bucket: TokenBucket): void {
+    if (bucket.tokens >= 1) {
+      bucket.tokens -= 1;
+    }
+  }
+
+  /**
+   * Start cleanup timer for expired sessions
+   */
+  private startCleanup(): void {
+    this.cleanupTimer = setInterval(() => {
+      const now = Date.now();
+      const expireTime = this.config.cleanupInterval * 2;
+
+      for (const [sessionId, bucket] of this.sessionBuckets) {
+        if (now - bucket.lastRefill > expireTime) {
+          this.sessionBuckets.delete(sessionId);
+          this.logger.debug('Rate limit session expired', { sessionId });
+        }
+      }
+    }, this.config.cleanupInterval);
+  }
+}
+
+export function createRateLimiter(
+  logger: ILogger,
+  config?: Partial<RateLimitConfig>
+): RateLimiter {
+  return new RateLimiter(logger, config);
+}
+
+/**
+ * Express/Connect middleware for rate limiting
+ */
+export function rateLimitMiddleware(rateLimiter: RateLimiter) {
+  return (req: any, res: any, next: () => void) => {
+    const sessionId = req.headers['x-session-id'] || req.ip;
+    const result = rateLimiter.check(sessionId);
+
+    res.setHeader('X-RateLimit-Remaining', result.remaining);
+    res.setHeader('X-RateLimit-Reset', Math.ceil(Date.now() / 1000) + Math.ceil(result.resetIn / 1000));
+
+    if (!result.allowed) {
+      res.setHeader('Retry-After', result.retryAfter);
+      res.status(429).json({
+        jsonrpc: '2.0',
+        id: null,
+        error: {
+          code: -32000,
+          message: 'Rate limit exceeded',
+          data: { retryAfter: result.retryAfter },
+        },
+      });
+      return;
+    }
+
+    rateLimiter.consume(sessionId);
+    next();
+  };
+}