@comake/skl-js-engine/src/JsExecutor/jsonRpc/JsonRpcServer.ts

Standard Knowledge Language Javascript Engine

import type {
  JsonRpcError,
  JsonRpcId,
  JsonRpcMessage,
  JsonRpcMethodHandler,
  JsonRpcMethodRegistry,
  JsonRpcNotification,
  JsonRpcRequest,
  JsonRpcResponse,
  JsonRpcServerConfig,
  PendingRequest,
  ValidationResult
} from './types';
import { JsonRpcErrorCode, JSONRPC_VERSION } from './types';

type TGetStatsResponse = {
  registeredMethods: number;
  pendingRequests: number;
  maxConcurrentRequests: number;
};

/**
 * Event emitter interface for JSON-RPC server events
 */
export interface JsonRpcServerEvents {
  request: (request: JsonRpcRequest) => void;
  notification: (notification: JsonRpcNotification) => void;
  response: (response: JsonRpcResponse) => void;
  error: (error: Error, request?: JsonRpcRequest) => void;
  timeout: (request: JsonRpcRequest) => void;
}

/**
 * JSON-RPC 2.0 Server implementation
 */
export class JsonRpcServer {
  private readonly methods: JsonRpcMethodRegistry = {};
  private readonly config: Required<JsonRpcServerConfig>;
  private readonly pendingRequests = new Map<JsonRpcId, PendingRequest>();
  private readonly listeners: Partial<JsonRpcServerEvents> = {};

  public constructor(config: JsonRpcServerConfig = {}) {
    this.config = {
      maxConcurrentRequests: config.maxConcurrentRequests ?? 100,
      requestTimeout: config.requestTimeout ?? 30000,
      strictValidation: config.strictValidation ?? true,
      errorHandler: config.errorHandler ?? this.defaultErrorHandler.bind(this)
    };
  }

  /**
   * Register a method handler
   * @param method - Method name
   * @param handler - Method handler function
   */
  public registerMethod<TParams, TResult>(method: string, handler: JsonRpcMethodHandler<TParams, TResult>): void {
    if (typeof method !== 'string' || method.length === 0) {
      throw new Error('Method name must be a non-empty string');
    }
    if (typeof handler !== 'function') {
      throw new Error('Handler must be a function');
    }
    this.methods[method] = handler as JsonRpcMethodHandler;
  }

  /**
   * Unregister a method handler
   * @param method - Method name
   */
  public unregisterMethod(method: string): void {
    Reflect.deleteProperty(this.methods, method);
  }

  /**
   * Get all registered methods
   * @returns Array of method names
   */
  public getRegisteredMethods(): string[] {
    return Object.keys(this.methods);
  }

  /**
   * Process an incoming JSON-RPC message
   * @param messageData - Raw message data (string or object)
   * @returns Promise resolving to response (if request) or undefined (if notification)
   */
  public async processMessage(messageData: string | object): Promise<JsonRpcResponse | undefined> {
    let message: JsonRpcMessage;

    try {
      // Parse message if it's a string
      if (typeof messageData === 'string') {
        try {
          message = JSON.parse(messageData);
        } catch {
          return this.createErrorResponse(null, {
            code: JsonRpcErrorCode.parseError,
            message: 'Parse error',
            data: 'Invalid JSON'
          });
        }
      } else {
        message = messageData as JsonRpcMessage;
      }

      // Validate message format
      const validation = this.validateMessage(message);
      if (!validation.valid) {
        return this.createErrorResponse(this.extractId(message), validation.error!);
      }

      // Handle request or notification
      if (this.isRequest(message)) {
        this.emit('request', message);
        return await this.handleRequest(message);
      }
      if (this.isNotification(message)) {
        this.emit('notification', message);
        await this.handleNotification(message);
        return;
      }
      // This shouldn't happen after validation, but just in case
      return this.createErrorResponse(this.extractId(message), {
        code: JsonRpcErrorCode.invalidRequest,
        message: 'Invalid request',
        data: 'Message is neither request nor notification'
      });
    } catch (error: unknown) {
      const jsonRpcError = this.config.errorHandler(error as Error, this.isRequest(message!) ? message : undefined);
      return this.createErrorResponse(this.extractId(message!), jsonRpcError);
    }
  }

  /**
   * Handle a JSON-RPC request
   * @param request - JSON-RPC request object
   * @returns Promise resolving to response
   */
  private async handleRequest(request: JsonRpcRequest): Promise<JsonRpcResponse> {
    // Check concurrent request limit
    if (this.pendingRequests.size >= this.config.maxConcurrentRequests) {
      return this.createErrorResponse(request.id, {
        code: JsonRpcErrorCode.internalError,
        message: 'Server overloaded',
        data: `Maximum concurrent requests (${this.config.maxConcurrentRequests}) exceeded`
      });
    }

    // Check if method exists
    const handler = this.methods[request.method];
    if (!handler) {
      return this.createErrorResponse(request.id, {
        code: JsonRpcErrorCode.methodNotFound,
        message: 'Method not found',
        data: `Method '${request.method}' is not registered`
      });
    }

    // Execute method with timeout
    const requestPromise = this.executeMethodWithTimeout(handler, request);

    try {
      const result = await requestPromise;
      return this.createSuccessResponse(request.id, result);
    } catch (error: unknown) {
      const jsonRpcError = this.config.errorHandler(error as Error, request);
      return this.createErrorResponse(request.id, jsonRpcError);
    } finally {
      this.pendingRequests.delete(request.id);
    }
  }

  /**
   * Handle a JSON-RPC notification
   * @param notification - JSON-RPC notification object
   */
  private async handleNotification(notification: JsonRpcNotification): Promise<void> {
    const handler = this.methods[notification.method];
    if (!handler) {
      // For notifications, we silently ignore unknown methods
      return;
    }

    try {
      // Handle both sync and async functions by wrapping in Promise.resolve
      const result = handler(notification.params);
      await Promise.resolve(result);
    } catch (error: unknown) {
      // For notifications, we emit error but don't send response
      this.emit('error', error as Error);
    }
  }

  /**
   * Execute method with timeout
   * @param handler - Method handler
   * @param request - JSON-RPC request
   * @returns Promise resolving to result
   */
  private async executeMethodWithTimeout<TResult>(
    handler: JsonRpcMethodHandler,
    request: JsonRpcRequest
  ): Promise<TResult> {
    return new Promise((resolve, reject): void => {
      const timeoutId = setTimeout((): void => {
        reject(new Error(`Method '${request.method}' timed out after ${this.config.requestTimeout}ms`));
      }, this.config.requestTimeout);

      try {
        const pendingRequest: PendingRequest<TResult> = {
          id: request.id,
          method: request.method,
          timestamp: Date.now(),
          timeout: this.config.requestTimeout,
          resolve(value: TResult) {
            clearTimeout(timeoutId);
            resolve(value);
          },
          reject(error: Error) {
            clearTimeout(timeoutId);
            reject(error);
          }
        };

        // Handle both sync and async functions by wrapping in Promise.resolve
        const result = handler(request.params, request.id);
        Promise.resolve(result)
          .then((finalResult: any): void => {
            clearTimeout(timeoutId);
            pendingRequest.resolve(finalResult);
          })
          .catch((error: unknown): void => {
            clearTimeout(timeoutId);
            pendingRequest.reject(error as Error);
          });
      } catch (error: unknown) {
        // Handle synchronous errors (like sync functions that throw)
        clearTimeout(timeoutId);
        reject(error);
      }
    });
  }

  /**
   * Create a pending request promise with timeout
   * @param request - JSON-RPC request
   * @param timeout - Timeout in milliseconds
   * @returns Promise resolving to the result
   */
  private async createPendingRequest<TResult>(request: JsonRpcRequest, timeout: number): Promise<TResult> {
    return new Promise<TResult>((resolve, reject) => {
      const timeoutId = setTimeout(() => {
        this.pendingRequests.delete(request.id);
        this.emit('timeout', request);
        reject(new Error(`Request ${request.id} timed out after ${timeout}ms`));
      }, timeout);

      const pendingRequest: PendingRequest<TResult> = {
        id: request.id,
        method: request.method,
        timestamp: Date.now(),
        timeout,
        resolve(value: TResult) {
          clearTimeout(timeoutId);
          resolve(value);
        },
        reject(error: Error) {
          clearTimeout(timeoutId);
          reject(error);
        }
      };

      this.pendingRequests.set(request.id, pendingRequest);
    });
  }

  /**
   * Validate JSON-RPC message format
   * @param message - Message to validate
   * @returns Validation result
   */
  private validateMessage(message: any): ValidationResult {
    if (!message || typeof message !== 'object') {
      return {
        valid: false,
        error: {
          code: JsonRpcErrorCode.invalidRequest,
          message: 'Invalid request',
          data: 'Message must be an object'
        }
      };
    }

    if (message.jsonrpc !== JSONRPC_VERSION) {
      return {
        valid: false,
        error: {
          code: JsonRpcErrorCode.invalidRequest,
          message: 'Invalid request',
          data: `Invalid jsonrpc version. Expected '${JSONRPC_VERSION}'`
        }
      };
    }

    if (typeof message.method !== 'string' || message.method.length === 0) {
      return {
        valid: false,
        error: {
          code: JsonRpcErrorCode.invalidRequest,
          message: 'Invalid request',
          data: 'Method must be a non-empty string'
        }
      };
    }

    // Check if it's a request (has id) or notification (no id)
    const hasId = 'id' in message;
    if (hasId && typeof message.id !== 'string' && typeof message.id !== 'number' && message.id !== null) {
      return {
        valid: false,
        error: {
          code: JsonRpcErrorCode.invalidRequest,
          message: 'Invalid request',
          data: 'ID must be a string, number, or null'
        }
      };
    }

    return { valid: true };
  }

  /**
   * Check if message is a request (has id)
   */
  private isRequest(message: any): message is JsonRpcRequest {
    return message && typeof message === 'object' && 'id' in message;
  }

  /**
   * Check if message is a notification (no id)
   */
  private isNotification(message: any): message is JsonRpcNotification {
    return message && typeof message === 'object' && !('id' in message);
  }

  /**
   * Extract ID from message safely
   */
  private extractId(message: any): JsonRpcId {
    if (message && typeof message === 'object' && 'id' in message) {
      return message.id;
    }
    return null;
  }

  /**
   * Create success response
   */
  private createSuccessResponse(id: JsonRpcId, result: any): JsonRpcResponse {
    return {
      jsonrpc: JSONRPC_VERSION,
      result,
      id
    };
  }

  /**
   * Create error response
   */
  private createErrorResponse(id: JsonRpcId, error: JsonRpcError): JsonRpcResponse {
    return {
      jsonrpc: JSONRPC_VERSION,
      error,
      id
    };
  }

  /**
   * Default error handler
   */
  private defaultErrorHandler(error: Error, request?: JsonRpcRequest): JsonRpcError {
    // Handle parameter validation errors
    if (error.message.includes('invalid param')) {
      return {
        code: JsonRpcErrorCode.invalidParams,
        message: 'Invalid params',
        data: error.message
      };
    }

    // Default to internal error
    return {
      code: JsonRpcErrorCode.internalError,
      message: error.message || 'Internal error',
      data: error.stack
    };
  }

  /**
   * Register event listener
   */
  public on<TKey extends keyof JsonRpcServerEvents>(event: TKey, listener: JsonRpcServerEvents[TKey]): void {
    this.listeners[event] = listener;
  }

  /**
   * Remove event listener
   */
  public off<TKey extends keyof JsonRpcServerEvents>(event: TKey, listener: JsonRpcServerEvents[TKey]): void {
    if (this.listeners[event] === listener) {
      Reflect.deleteProperty(this.listeners, event);
    }
  }

  /**
   * Emit event
   */
  private emit<TKey extends keyof JsonRpcServerEvents>(
    event: TKey,
    ...args: Parameters<JsonRpcServerEvents[TKey]>
  ): void {
    const listener = this.listeners[event];
    if (listener) {
      (listener as any)(...args);
    }
  }

  /**
   * Get server statistics
   */
  public getStats(): TGetStatsResponse {
    return {
      registeredMethods: Object.keys(this.methods).length,
      pendingRequests: this.pendingRequests.size,
      maxConcurrentRequests: this.config.maxConcurrentRequests
    };
  }

  /**
   * Shutdown server (reject all pending requests)
   */
  public async shutdown(): Promise<void> {
    // Reject all pending requests
    for (const [ id, pendingRequest ] of this.pendingRequests) {
      try {
        pendingRequest.reject(new Error('Server shutdown'));
      } catch {
        // Ignore errors during shutdown
      }
    }

    this.pendingRequests.clear();
  }
}