@comake/skl-js-engine/src/JsExecutor/jsExecutor.ts

Standard Knowledge Language Javascript Engine

/* eslint-disable @typescript-eslint/naming-convention */
/* eslint-disable callback-return */

/* eslint-disable func-style */

import { NotInitializedError } from './errors';
import type { TransportConfig } from './transport';
import { JsonRpcServer, ParentStdioTransport, STANDARD_METHODS } from './transport';
import type { ExecutionOptions, IShutdownOptions } from './types';

export type IJSCallbacks = Record<string, (...args: any[]) => any>;

/**
 * JSExecutor V2 with enhanced abort control for immediate shutdown
 *
 * @example Writing abort-aware callbacks:
 * ```typescript
 * const jsExecutor = new JSExecutorV2(scriptPath, {
 *   async myCallback(): Promise<string> {
 *     return new Promise((resolve, reject) => {
 *       const timeoutId = setTimeout(() => resolve('done'), 5000);
 *
 *       // Listen to the global abort signal to clean up resources
 *       const abortSignal = JSExecutorV2.getGlobalAbortSignal();
 *       const onAbort = () => {
 *         clearTimeout(timeoutId);
 *         reject(new Error('Operation aborted'));
 *       };
 *
 *       if (abortSignal.aborted) {
 *         onAbort();
 *         return;
 *       }
 *
 *       abortSignal.addEventListener('abort', onAbort);
 *     });
 *   }
 * });
 * ```
 */
export class JSExecutor {
  private isInitialized = false;
  private server?: JsonRpcServer;
  private transport?: ParentStdioTransport;
  private readonly abortController = new AbortController();
  private static readonly globalAbortController = new AbortController();
  private shutdownTimeout?: NodeJS.Timeout;

  public constructor(
    private readonly executorScriptPath: string,
    private readonly callbacks: Record<string, (...args: any[]) => any>
  ) {}

  /**
   * Static method to check if any JSExecutor instance is being shut down
   * Callbacks can use this to check if they should abort their operations
   */
  public static isShuttingDown(): boolean {
    return JSExecutor.globalAbortController.signal.aborted;
  }

  /**
   * Get the global abort signal that callbacks can listen to
   * This allows callbacks to immediately abort when any JSExecutor shuts down
   */
  public static getGlobalAbortSignal(): AbortSignal {
    return JSExecutor.globalAbortController.signal;
  }

  public async initialize(config?: TransportConfig, executionOptions?: ExecutionOptions): Promise<void> {
    if (this.isInitialized) {
      return;
    }

    this.server = new JsonRpcServer({
      // 1 hour
      requestTimeout: 60 * 60 * 1000
    });

    // Register callbacks with the server
    Object.entries(this.callbacks).forEach(([ methodName, callback ]) => {
      this.server!.registerMethod(methodName, this.wrapCallbackWithAbort(callback));
    });

    // Create and initialize the transport once
    this.transport = new ParentStdioTransport(this.executorScriptPath, this.server);

    await this.transport.initialize(config, executionOptions);

    this.isInitialized = true;
  }

  public async shutdown(options?: IShutdownOptions): Promise<void> {
    // Clear any existing shutdown timeout
    if (this.shutdownTimeout) {
      clearTimeout(this.shutdownTimeout);
      this.shutdownTimeout = undefined;
    }

    if (options?.afterTimeout) {
      this.shutdownTimeout = setTimeout(async() => {
        await this.shutdownCore();
      }, options.afterTimeout);
    } else {
      await this.shutdownCore();
    }
  }

  private async shutdownCore(): Promise<void> {
    // // Signal global abort first
    // JSExecutor.globalAbortController.abort();

    // Then signal this instance's abort
    this.abortController.abort();

    if (this.transport) {
      await this.transport.close();
      this.transport = undefined;
    }

    this.isInitialized = false;
    this.server = undefined;
  }

  private async rpcComplaintCallback(callback: (...args: any[]) => any, args: any[]): Promise<any> {
    const result = await callback(...args);
    // RPC expects a result object
    return result ?? {};
  }

  /**
   * Wrap a callback with abort functionality
   * The callback will be automatically aborted if JSExecutor is shut down
   */
  private wrapCallbackWithAbort<T extends (...args: any[]) => any>(callback: T): T {
    return ((...args: any[]) => {
      // Check if we're shutting down before executing
      if (this.abortController.signal.aborted || JSExecutor.globalAbortController.signal.aborted) {
        return Promise.reject(new Error('JSExecutor is shutting down'));
      }

      const callbackPromise = Promise.resolve(this.rpcComplaintCallback(callback, args));

      // Create an abort promise
      const abortPromise = new Promise((_, reject) => {
        const onAbort = (): void => {
          reject(new Error('JSExecutor is shutting down'));
        };

        // Listen to both instance and global abort signals
        this.abortController.signal.addEventListener('abort', onAbort, { once: true });
        JSExecutor.globalAbortController.signal.addEventListener('abort', onAbort, { once: true });
      });

      // Race the callback against the abort signal
      return Promise.race([ callbackPromise, abortPromise ]);
    }) as T;
  }

  public async execute(code: string, args: Record<string, any>, executionOptions: ExecutionOptions): Promise<any> {
    if (!this.isInitialized) {
      throw new NotInitializedError('JSExecutor is not initialized, please call initialize() first');
    }

    if (!this.transport || !this.transport.isReady()) {
      throw new Error('Transport is not ready');
    }

    const response = await this.transport.request(
      STANDARD_METHODS.executeCode,
      {
        code,
        args,
        functionName: executionOptions.functionName
      },
      {
        timeout: executionOptions.timeout,
        retries: executionOptions.retries ?? 0
      }
    );
    return response;
  }

  /**
   * Legacy method for backward compatibility
   * @deprecated Use execute() instead
   */
  public async call(methodName: string, ...args: any[]): Promise<any> {
    if (!this.transport || !this.transport.isReady()) {
      throw new Error('Transport is not ready');
    }
    return this.transport.request(methodName, args);
  }
}