@comake/skl-js-engine/src/JsExecutor/transport/stdio/ChildStdioTransport.ts

Standard Knowledge Language Javascript Engine

/* eslint-disable indent */
import { EventEmitter } from 'node:events';
import type { ClientTransport } from '../../jsonRpc/JsonRpcClient';
import { JsonRpcClient } from '../../jsonRpc/JsonRpcClient';
import { JsonRpcServer } from '../../jsonRpc/JsonRpcServer';
import { isResponse, MessageBuffer } from '../utils/MessageUtils';

/**
 * Custom logger that outputs to stderr to avoid contaminating stdout JSON-RPC channel
 */
class ChildProcessLogger {
  private readonly name?: string;
  private metadata: Record<string, any> = {};

  public setMetadata(metadata: Record<string, any>): void {
    this.metadata = { ...this.metadata, ...metadata };
  }

  public log(message: string, data?: any): void {
    const logEntry = {
      message,
      metadata: this.metadata,
      ...data && { data }
    };
    // Use stderr to avoid contaminating stdout JSON-RPC channel
    process.stderr.write(`${message} ${JSON.stringify(logEntry.metadata)} ${data ? JSON.stringify(data) : ''}\n`);
  }

  public error(message: string, error?: Error): void {
    const logEntry = {
      level: 'error',
      message,
      metadata: this.metadata,
      ...error && { error: error.message, stack: error.stack }
    };
    process.stderr.write(`ERROR: ${JSON.stringify(logEntry)}\n`);
  }
}

/**
 * Bidirectional StdioTransport for child processes
 * This class handles both server (receiving requests) and client (sending requests) functionality
 * over a single stdio channel with proper message routing.
 *
 * Usage in child process:
 * ```typescript
 * const transport = new ChildStdioTransport();
 *
 * // Register methods that parent can call
 * transport.registerMethod('ping', async () => 'pong');
 *
 * // Make requests to parent
 * const result = await transport.request('getTime');
 * ```
 */
export class ChildStdioTransport extends EventEmitter {
  private readonly server: JsonRpcServer;
  private readonly client: JsonRpcClient;
  private readonly logger = new ChildProcessLogger();
  private readonly messageBuffer = new MessageBuffer();
  private name?: string;
  private initialized = false;

  public constructor() {
    super();
    this.server = new JsonRpcServer({
      requestTimeout: 60 * 60 * 1000
    });
    this.client = new JsonRpcClient({
      defaultTimeout: 60 * 60 * 1000
    });

    // Set up the client transport that routes through our stdio
    const clientTransport: ClientTransport = {
      send: async(message: string): Promise<void> => {
        this.sendMessage(message);
      },
      onMessage(): void {
        // The message routing will handle client responses automatically
      },
      async close(): Promise<void> {
        // Nothing to close for stdio
      }
    };

    this.client.setTransport(clientTransport);
    this.setupEventHandlers();
  }

  /**
   * Set a name for this transport (used in logging)
   */
  public setName(name: string): void {
    this.name = name;
    this.logger.setMetadata({ name, transport: 'ChildStdioTransport' });
  }

  /**
   * Initialize the bidirectional transport
   * This should be called once after setting up all methods
   */
  public async initialize(): Promise<void> {
    if (this.initialized) {
      throw new Error('Transport already initialized');
    }

    this.initialized = true;
    this.setupStdioCommunication();
    this.logger.log('Transport initialized');
  }

  /**
   * Register a method that the parent can call
   * @param method - Method name
   * @param handler - Method handler function
   */
  public registerMethod<TParams = any, TResult = any>(
    method: string,
    handler: (params: TParams) => Promise<TResult> | TResult
  ): void {
    this.server.registerMethod(method, handler);
  }

  /**
   * Send a request to the parent process
   * @param method - Method name
   * @param params - Method parameters
   * @param options - Request options
   * @returns Promise resolving to the response
   */
  public async request<TParams = any, TResult = any>(
    method: string,
    params?: TParams,
    options?: { timeout?: number }
  ): Promise<TResult> {
    if (!this.initialized) {
      throw new Error('Transport not initialized. Call initialize() first.');
    }
    return this.client.request<TParams, TResult>(method, params, options);
  }

  /**
   * Send a notification to the parent process (no response expected)
   * @param method - Method name
   * @param params - Method parameters
   */
  public async notify<TParams = any>(method: string, params?: TParams): Promise<void> {
    if (!this.initialized) {
      throw new Error('Transport not initialized. Call initialize() first.');
    }
    return this.client.notify(method, params);
  }

  /**
   * Get transport statistics
   */
  public getStats(): {
    serverMethods: number;
    pendingRequests: number;
    initialized: boolean;
  } {
    return {
      serverMethods: 0,
      pendingRequests: this.client.getStats().pendingRequests,
      initialized: this.initialized
    };
  }

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

  /**
   * Set up stdio communication handlers
   */
  private setupStdioCommunication(): void {
    // Handle incoming data from parent
    process.stdin.on('data', async(data: Buffer): Promise<void> => {
      const messages = this.messageBuffer.addData(data.toString());

      for (const messageData of messages) {
        await this.handleIncomingMessage(messageData);
      }
    });

    // Handle process termination gracefully
    process.on('SIGTERM', (): void => {
      this.close().finally((): void => process.exit(0));
    });

    process.on('SIGINT', (): void => {
      this.close().finally((): void => process.exit(0));
    });
  }

  /**
   * Handle incoming message and route to appropriate handler
   */
  private async handleIncomingMessage(messageData: string): Promise<void> {
    try {
      const message = JSON.parse(messageData);

      // Check if this is a response to our request (has 'result' or 'error' and 'id')
      if (isResponse(message)) {
        // This is a response to a request we made - route to client
        this.logger.log(`Routing response to client: ${message.id}`);
        await this.client.handleIncomingMessage(messageData);
      } else {
        // This is a request or notification for us - route to server
        this.logger.log(`Routing request to server: ${message.method}`);
        const response = await this.server.processMessage(messageData);
        this.logger.log('Sending response to parent process', response);
        // Only send response if there is one (requests get responses, notifications don't)
        if (response) {
          this.sendMessage(JSON.stringify(response));
        }
      }
    } catch (error: unknown) {
      process.stderr.write(`Error processing message: ${(error as Error).message}\n`);
      this.emit('error', error);
    }
  }

  /**
   * Send a message to the parent process
   */
  private sendMessage(message: string): void {
    this.logger.log(`Sending: ${message}`);
    process.stdout.write(`${message}\n`);
  }

  /**
   * Set up event handlers
   */
  private setupEventHandlers(): void {
    this.server.on('error', (error: Error): void => {
      this.logger.error('Server error', error);
      this.emit('error', error);
    });

    this.client.on('error', (error: Error): void => {
      this.logger.error('Client error', error);
      this.emit('error', error);
    });
  }
}