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

Standard Knowledge Language Javascript Engine

/* eslint-disable indent */

import type {
  JsonRpcClientConfig,
  JsonRpcId,
  JsonRpcNotification,
  JsonRpcParams,
  JsonRpcRequest,
  JsonRpcResponse,
  JsonRpcSuccessResponse,
  PendingRequest
} from './types';
import { JsonRpcErrorCode, JSONRPC_VERSION } from './types';

/**
 * Event emitter interface for JSON-RPC client events
 */
export interface JsonRpcClientEvents {
  request: (request: JsonRpcRequest) => void;
  response: (response: JsonRpcResponse) => void;
  error: (error: Error) => void;
  timeout: (request: JsonRpcRequest) => void;
}

/**
 * Transport interface for sending messages
 */
export interface ClientTransport {
  send: (message: string) => Promise<void>;
  onMessage: (handler: (message: string) => void) => void;
  close: () => Promise<void>;
}

/**
 * JSON-RPC 2.0 Client implementation
 */
export class JsonRpcClient {
  private readonly config: Required<JsonRpcClientConfig>;
  private readonly pendingRequests = new Map<JsonRpcId, PendingRequest>();
  private readonly listeners: Partial<JsonRpcClientEvents> = {};
  private requestIdCounter = 0;
  private transport?: ClientTransport;

  public constructor(config: JsonRpcClientConfig = {}) {
    this.config = {
      defaultTimeout: config.defaultTimeout ?? 30000,
      maxRetries: config.maxRetries ?? 3,
      retryDelay: config.retryDelay ?? 1000,
      sequentialIds: config.sequentialIds ?? true
    };
  }

  /**
   * Set the transport for sending messages
   * @param transport - Transport implementation
   */
  public setTransport(transport: ClientTransport): void {
    this.transport = transport;
    this.transport.onMessage((message: string): void => {
      this.handleIncomingMessage(message).catch((error): void => {
        this.emit('error', error);
      });
    });
  }

  /**
   * Send a JSON-RPC request and wait for response
   * @param method - Method name
   * @param params - Method parameters
   * @param options - Request options
   * @returns Promise resolving to the response result
   */
  public async request<TParams = JsonRpcParams, TResult = any>(
    method: string,
    params?: TParams,
    options: { timeout?: number; retries?: number } = {}
  ): Promise<TResult> {
    if (!this.transport) {
      throw new Error('Transport not set. Call setTransport() first.');
    }

    const timeout = options.timeout ?? this.config.defaultTimeout;
    const maxRetries = options.retries ?? this.config.maxRetries;

    let lastError: Error | undefined;

    for (let attempt = 0; attempt <= maxRetries; attempt++) {
      try {
        const request = this.createRequest(method, params);
        const result = await this.sendRequestWithTimeout<TResult>(request, timeout);
        return result;
      } catch (error: unknown) {
        lastError = error as Error;
        // Don't retry on certain errors
        if (this.shouldNotRetry(error as Error) || attempt === maxRetries) {
          break;
        }

        // Wait before retrying
        if (attempt < maxRetries) {
          await this.delay(this.config.retryDelay * (attempt + 1));
        }
      }
    }

    throw lastError ?? new Error('Request failed after all retry attempts');
  }

  /**
   * Send a JSON-RPC notification (no response expected)
   * @param method - Method name
   * @param params - Method parameters
   */
  public async notify<TParams = JsonRpcParams>(method: string, params?: TParams): Promise<void> {
    if (!this.transport) {
      throw new Error('Transport not set. Call setTransport() first.');
    }

    const notification = this.createNotification(method, params);
    await this.transport.send(JSON.stringify(notification));
  }

  /**
   * Send a batch of requests
   * @param requests - Array of request specifications
   * @returns Promise resolving to array of results
   */
  public async batch(requests: { method: string; params?: any; id?: JsonRpcId }[]): Promise<any[]> {
    if (!this.transport) {
      throw new Error('Transport not set. Call setTransport() first.');
    }

    const batchRequests = requests.map((req): JsonRpcRequest => {
      const id = req.id ?? this.generateId();
      return this.createRequest(req.method, req.params, id);
    });

    const batchMessage = JSON.stringify(batchRequests);

    // Create pending promises for all requests
    const promises = batchRequests.map(
      (request): Promise<any> => this.createPendingRequest<any>(request, this.config.defaultTimeout)
    );

    // Send the batch
    await this.transport.send(batchMessage);

    // Wait for all responses
    return Promise.all(promises);
  }

  /**
   * Handle incoming message from transport
   * @param messageData - Raw message data
   */
  public async handleIncomingMessage(messageData: string): Promise<void> {
    try {
      const message = JSON.parse(messageData);
      // Handle batch response
      if (Array.isArray(message)) {
        for (const response of message) {
          this.handleResponse(response);
        }
      } else {
        this.handleResponse(message);
      }
    } catch (error: unknown) {
      this.emit('error', new Error(`Failed to parse incoming message: ${(error as Error).message}`));
    }
  }

  /**
   * Handle a single response message
   * @param response - JSON-RPC response
   */
  private handleResponse(response: JsonRpcResponse): void {
    this.emit('response', response);

    const pendingRequest = this.pendingRequests.get(response.id);
    if (!pendingRequest) {
      // Orphaned response - might be from a timed-out request
      return;
    }

    this.pendingRequests.delete(response.id);

    if (this.isSuccessResponse(response)) {
      pendingRequest.resolve(response.result);
    } else {
      const error = new Error(response.error.message);
      (error as any).code = response.error.code;
      (error as any).data = response.error.data;
      pendingRequest.reject(error);
    }
  }

  /**
   * Send a request with timeout handling
   * @param request - JSON-RPC request
   * @param timeout - Timeout in milliseconds
   * @returns Promise resolving to the result
   */
  private async sendRequestWithTimeout<TResult>(request: JsonRpcRequest, timeout: number): Promise<TResult> {
    const promise = this.createPendingRequest<TResult>(request, timeout);

    this.emit('request', request);
    await this.transport!.send(JSON.stringify(request));

    return promise;
  }

  /**
   * 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);
    });
  }

  /**
   * Create a JSON-RPC request
   * @param method - Method name
   * @param params - Method parameters
   * @param id - Optional request ID
   * @returns JSON-RPC request object
   */
  private createRequest<TParams = JsonRpcParams>(
    method: string,
    params?: TParams,
    id?: JsonRpcId
  ): JsonRpcRequest<TParams> {
    const request: JsonRpcRequest<TParams> = {
      jsonrpc: JSONRPC_VERSION,
      method,
      id: id ?? this.generateId()
    };

    if (params !== undefined) {
      request.params = params;
    }

    return request;
  }

  /**
   * Create a JSON-RPC notification
   * @param method - Method name
   * @param params - Method parameters
   * @returns JSON-RPC notification object
   */
  private createNotification<TParams = JsonRpcParams>(method: string, params?: TParams): JsonRpcNotification<TParams> {
    const notification: JsonRpcNotification<TParams> = {
      jsonrpc: JSONRPC_VERSION,
      method
    };

    if (params !== undefined) {
      notification.params = params;
    }

    return notification;
  }

  /**
   * Generate a unique request ID
   * @returns Unique request ID
   */
  private generateId(): JsonRpcId {
    if (this.config.sequentialIds) {
      this.requestIdCounter += 1;
      return this.requestIdCounter;
    }
    return `${Date.now()}-${Math.random().toString(36).slice(2, 11)}`;
  }

  /**
   * Check if response is a success response
   * @param response - JSON-RPC response
   * @returns True if success response
   */
  private isSuccessResponse(response: JsonRpcResponse): response is JsonRpcSuccessResponse {
    return 'result' in response;
  }

  /**
   * Check if an error should not be retried
   * @param error - Error to check
   * @returns True if error should not be retried
   */
  private shouldNotRetry(error: Error): boolean {
    // Don't retry on method not found, invalid params, or parse errors
    const errorCode = (error as any).code;
    return (
      errorCode === JsonRpcErrorCode.methodNotFound ||
      errorCode === JsonRpcErrorCode.invalidParams ||
      errorCode === JsonRpcErrorCode.parseError ||
      errorCode === JsonRpcErrorCode.invalidRequest ||
      errorCode === JsonRpcErrorCode.internalError
    );
  }

  /**
   * Delay execution for specified milliseconds
   * @param ms - Delay in milliseconds
   * @returns Promise that resolves after delay
   */
  private async delay(ms: number): Promise<void> {
    return new Promise(resolve => setTimeout(resolve, ms));
  }

  /**
   * Register event listener
   * @param event - Event name
   * @param listener - Event listener
   */
  public on<TKey extends keyof JsonRpcClientEvents>(event: TKey, listener: JsonRpcClientEvents[TKey]): void {
    this.listeners[event] = listener;
  }

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

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

  /**
   * Get client statistics
   * @returns Client statistics
   */
  public getStats(): {
    pendingRequests: number;
    totalRequests: number;
    config: Required<JsonRpcClientConfig>;
  } {
    return {
      pendingRequests: this.pendingRequests.size,
      totalRequests: this.requestIdCounter,
      config: { ...this.config }
    };
  }

  /**
   * Cancel all pending requests
   */
  public cancelAllRequests(): void {
    const error = new Error('Request cancelled');
    for (const [ id, pendingRequest ] of this.pendingRequests) {
      pendingRequest.reject(error);
    }
    this.pendingRequests.clear();
  }

  /**
   * Cancel a specific request
   * @param id - Request ID to cancel
   * @returns True if request was cancelled
   */
  public cancelRequest(id: JsonRpcId): boolean {
    const pendingRequest = this.pendingRequests.get(id);
    if (pendingRequest) {
      pendingRequest.reject(new Error('Request cancelled'));
      this.pendingRequests.delete(id);
      return true;
    }
    return false;
  }

  /**
   * Close the client and cleanup resources
   */
  public async close(): Promise<void> {
    this.cancelAllRequests();
    if (this.transport) {
      await this.transport.close();
    }
  }
}