tywrap 0.2.1 → 0.4.0

package/src/runtime/node.ts

@@ -18,10 +18,11 @@ import { getVenvBinDir, getVenvPythonExe } from '../utils/runtime.js';
 import { globalCache } from '../utils/cache.js';
 
 import { BridgeProtocol, type BridgeProtocolOptions } from './bridge-protocol.js';
-import { BridgeProtocolError } from './errors.js';
+import { BridgeExecutionError, BridgeProtocolError } from './errors.js';
 import { ProcessIO } from './process-io.js';
 import { PooledTransport } from './pooled-transport.js';
 import type { CodecOptions } from './safe-codec.js';
+import { PROTOCOL_ID } from './transport.js';
 import type { PooledWorker } from './worker-pool.js';
 
 // =============================================================================
@@ -73,8 +74,7 @@ export interface NodeBridgeOptions {
 
   /** Commands to run on each process at startup for warming up. */
   warmupCommands?: Array<
-    | { module: string; functionName: string; args?: unknown[] }
-    | { method: string; params: unknown } // Legacy format for backwards compatibility
+    { module: string; functionName: string; args?: unknown[] } | { method: string; params: unknown } // Legacy shape preserved so runtime can surface a migration error
   >;
 
   // ===========================================================================
@@ -120,12 +120,17 @@ interface ResolvedOptions {
   enableCache: boolean;
   env: Record<string, string | undefined>;
   codec?: CodecOptions;
-  warmupCommands: Array<
-    | { module: string; functionName: string; args?: unknown[] }
-    | { method: string; params: unknown }
-  >;
+  warmupCommands: WarmupCommand[];
+}
+
+interface WarmupCommand {
+  module: string;
+  functionName: string;
+  args?: unknown[];
 }
 
+const WORKER_READY_TIMEOUT_MS = 5000;
+
 // =============================================================================
 // UTILITIES
 // =============================================================================
@@ -158,6 +163,10 @@ function resolveVirtualEnv(
  * Get the environment variable key for PATH (case-insensitive on Windows).
  */
 function getPathKey(env: Record<string, string | undefined>): string {
+  if (Object.prototype.hasOwnProperty.call(env, 'PATH')) {
+    return 'PATH';
+  }
+
   for (const key of Object.keys(env)) {
     if (key.toLowerCase() === 'path') {
       return key;
@@ -166,6 +175,93 @@ function getPathKey(env: Record<string, string | undefined>): string {
   return 'PATH';
 }
 
+function setPathValue(env: Record<string, string>, value: string): void {
+  setEnvValue(env, 'PATH', value);
+
+  if (process.platform !== 'win32') {
+    return;
+  }
+
+  for (const key of Object.keys(env)) {
+    if (key !== 'PATH' && key.toLowerCase() === 'path') {
+      setEnvValue(env, key, value);
+    }
+  }
+}
+
+const DANGEROUS_ENV_OVERRIDE_KEYS = new Set(['__proto__', 'prototype', 'constructor']);
+
+function createNullPrototypeEnv(): Record<string, string> {
+  return Object.create(null) as Record<string, string>;
+}
+
+function setEnvValue(env: Record<string, string>, key: string, value: string): void {
+  Object.defineProperty(env, key, {
+    value,
+    writable: true,
+    enumerable: true,
+    configurable: true,
+  });
+}
+
+function getEnvValue(env: Record<string, string>, key: string): string | undefined {
+  const value = Reflect.get(env, key);
+  return typeof value === 'string' ? value : undefined;
+}
+
+function assertSafeEnvOverrideKey(key: string): void {
+  if (DANGEROUS_ENV_OVERRIDE_KEYS.has(key)) {
+    throw new BridgeProtocolError(`Invalid environment override key "${key}" in options.env`);
+  }
+}
+
+function normalizeWarmupCommands(commands: NodeBridgeOptions['warmupCommands']): WarmupCommand[] {
+  if (commands === undefined) {
+    return [];
+  }
+  if (!Array.isArray(commands)) {
+    throw new BridgeProtocolError('warmupCommands must be an array when provided');
+  }
+
+  const warmups = commands;
+  return warmups.map((command, index) => {
+    if (!command || typeof command !== 'object' || Array.isArray(command)) {
+      throw new BridgeProtocolError(
+        `Invalid warmup command at index ${index + 1}: expected { module, functionName, args? }`
+      );
+    }
+
+    const candidate = command as Record<string, unknown>;
+    if ('method' in candidate || 'params' in candidate) {
+      throw new BridgeProtocolError(
+        `Invalid warmup command at index ${index + 1}: legacy { method, params } format is no longer supported. Use { module, functionName, args? }.`
+      );
+    }
+
+    if (typeof candidate.module !== 'string' || candidate.module.trim().length === 0) {
+      throw new BridgeProtocolError(
+        `Invalid warmup command at index ${index + 1}: "module" must be a non-empty string`
+      );
+    }
+    if (typeof candidate.functionName !== 'string' || candidate.functionName.trim().length === 0) {
+      throw new BridgeProtocolError(
+        `Invalid warmup command at index ${index + 1}: "functionName" must be a non-empty string`
+      );
+    }
+    if (candidate.args !== undefined && !Array.isArray(candidate.args)) {
+      throw new BridgeProtocolError(
+        `Invalid warmup command at index ${index + 1}: "args" must be an array when provided`
+      );
+    }
+
+    return {
+      module: candidate.module,
+      functionName: candidate.functionName,
+      args: candidate.args as unknown[] | undefined,
+    };
+  });
+}
+
 // =============================================================================
 // NODE BRIDGE
 // =============================================================================
@@ -226,6 +322,8 @@ export class NodeBridge extends BridgeProtocol {
     const maxProcesses = options.maxProcesses ?? 1;
     const minProcesses = Math.min(options.minProcesses ?? 1, maxProcesses);
 
+    const warmupCommands = normalizeWarmupCommands(options.warmupCommands);
+
     const resolvedOptions: ResolvedOptions = {
       minProcesses,
       maxProcesses,
@@ -240,20 +338,20 @@ export class NodeBridge extends BridgeProtocol {
       enableCache: options.enableCache ?? false,
       env: options.env ?? {},
       codec: options.codec,
-      warmupCommands: options.warmupCommands ?? [],
+      warmupCommands,
     };
-
     // Build environment for ProcessIO
     const processEnv = buildProcessEnv(resolvedOptions);
 
-    // Create warmup callback for per-worker initialization
-    const onWorkerReady = resolvedOptions.warmupCommands.length > 0
-      ? createWarmupCallback(resolvedOptions.warmupCommands, resolvedOptions.timeoutMs)
-      : undefined;
+    const userWarmup =
+      resolvedOptions.warmupCommands.length > 0
+        ? createWarmupCallback(resolvedOptions.warmupCommands, resolvedOptions.timeoutMs)
+        : undefined;
+    const onWorkerReady = createWorkerReadyCallback(resolvedOptions.timeoutMs, userWarmup);
 
     // Create pooled transport with ProcessIO workers
     const transport = new PooledTransport({
-      createTransport: () =>
+      createTransport: (): ProcessIO =>
         new ProcessIO({
           pythonPath: resolvedOptions.pythonPath,
           bridgeScript: resolvedOptions.scriptPath,
@@ -265,6 +363,7 @@ export class NodeBridge extends BridgeProtocol {
       queueTimeoutMs: resolvedOptions.queueTimeoutMs,
       maxConcurrentPerWorker: resolvedOptions.maxConcurrentPerProcess,
       onWorkerReady,
+      onReplacementWorkerReady: onWorkerReady,
     });
 
     // Initialize BridgeProtocol with pooled transport
@@ -460,8 +559,74 @@ export class NodeBridge extends BridgeProtocol {
  * Simple request ID generator for warmup commands.
  */
 let warmupRequestId = 0;
-function generateWarmupId(): string {
-  return `warmup_${++warmupRequestId}_${Date.now()}`;
+function generateWarmupId(): number {
+  return ++warmupRequestId;
+}
+
+async function sendWarmupRequest(
+  worker: PooledWorker,
+  timeoutMs: number,
+  requestId: number,
+  label: string,
+  messagePayload: Record<string, unknown>
+): Promise<void> {
+  let message: string;
+  try {
+    message = JSON.stringify({
+      id: requestId,
+      protocol: PROTOCOL_ID,
+      ...messagePayload,
+    });
+  } catch (error) {
+    throw new BridgeExecutionError(
+      `${label} failed to encode request: ${error instanceof Error ? error.message : String(error)}`,
+      { cause: error instanceof Error ? error : undefined }
+    );
+  }
+
+  let response: string;
+  try {
+    response = await worker.transport.send(message, timeoutMs);
+  } catch (error) {
+    throw new BridgeExecutionError(
+      `${label} failed to send: ${error instanceof Error ? error.message : String(error)}`,
+      { cause: error instanceof Error ? error : undefined }
+    );
+  }
+
+  let parsed: unknown;
+  try {
+    parsed = JSON.parse(response);
+  } catch (error) {
+    throw new BridgeExecutionError(`${label} returned invalid JSON response`, {
+      cause: error instanceof Error ? error : undefined,
+    });
+  }
+
+  if (!parsed || typeof parsed !== 'object' || Array.isArray(parsed)) {
+    throw new BridgeExecutionError(
+      `${label} returned malformed response envelope for request ${requestId}`
+    );
+  }
+
+  const envelope = parsed as { result?: unknown; error?: unknown };
+  if ('error' in envelope && envelope.error !== undefined && envelope.error !== null) {
+    const errorPayload = envelope.error;
+    if (errorPayload && typeof errorPayload === 'object' && !Array.isArray(errorPayload)) {
+      const err = errorPayload as { type?: unknown; message?: unknown };
+      const errType = typeof err.type === 'string' ? err.type : 'Error';
+      const errMessage = typeof err.message === 'string' ? err.message : 'Unknown warmup error';
+      throw new BridgeExecutionError(`${label} failed: ${errType}: ${errMessage}`);
+    }
+
+    throw new BridgeExecutionError(`${label} failed with malformed error payload`);
+  }
+
+  if (!('result' in envelope)) {
+    throw new BridgeExecutionError(
+      `${label} returned malformed response envelope for request ${requestId}`
+    );
+  }
 }
 
 /**
@@ -471,38 +636,138 @@ function generateWarmupId(): string {
  * bypassing the pool to ensure each worker gets warmed up individually.
  */
 function createWarmupCallback(
-  warmupCommands: Array<
-    | { module: string; functionName: string; args?: unknown[] }
-    | { method: string; params: unknown }
-  >,
+  warmupCommands: WarmupCommand[],
   timeoutMs: number
 ): (worker: PooledWorker) => Promise<void> {
   return async (worker: PooledWorker) => {
-    for (const cmd of warmupCommands) {
-      try {
-        // Handle both new and legacy warmup command formats
-        if ('module' in cmd && 'functionName' in cmd) {
-          // Build the protocol message
-          const message = JSON.stringify({
-            id: generateWarmupId(),
-            type: 'call',
-            module: cmd.module,
-            functionName: cmd.functionName,
-            args: cmd.args ?? [],
-            kwargs: {},
-          });
-
-          // Send directly to this worker's transport
-          await worker.transport.send(message, timeoutMs);
+    for (const [index, cmd] of warmupCommands.entries()) {
+      const commandLabel = `${cmd.module}.${cmd.functionName}`;
+      await executeWorkerCall(
+        worker,
+        {
+          module: cmd.module,
+          functionName: cmd.functionName,
+          args: cmd.args ?? [],
+        },
+        timeoutMs,
+        {
+          label: `Warmup command #${index + 1} (${commandLabel})`,
+          invalidJsonMessage: 'returned invalid JSON response',
+          malformedEnvelopeMessage: requestId =>
+            `returned malformed response envelope for request ${requestId}`,
+          pythonErrorMessage: (errorType, errorMessage) => `failed: ${errorType}: ${errorMessage}`,
+          malformedErrorPayloadMessage: 'failed with malformed error payload',
         }
-        // Legacy format { method, params } is ignored as it's not supported
-      } catch {
-        // Ignore warmup errors - they're not critical
-      }
+      );
     }
   };
 }
 
+function createWorkerReadyCallback(
+  timeoutMs: number,
+  extraWarmup?: (worker: PooledWorker) => Promise<void>
+): (worker: PooledWorker) => Promise<void> {
+  return async (worker: PooledWorker) => {
+    const readyTimeoutMs = timeoutMs > 0 ? Math.max(timeoutMs, WORKER_READY_TIMEOUT_MS) : 0;
+    await sendWarmupRequest(worker, readyTimeoutMs, generateWarmupId(), 'Worker warmup check', {
+      method: 'meta',
+      params: {},
+    });
+
+    await extraWarmup?.(worker);
+  };
+}
+
+interface WorkerCallSpec {
+  module: string;
+  functionName: string;
+  args: unknown[];
+}
+
+interface WorkerCallErrorMessages {
+  label: string;
+  invalidJsonMessage: string;
+  malformedEnvelopeMessage: (requestId: number) => string;
+  pythonErrorMessage: (errorType: string, errorMessage: string) => string;
+  malformedErrorPayloadMessage: string;
+}
+
+async function executeWorkerCall(
+  worker: PooledWorker,
+  request: WorkerCallSpec,
+  timeoutMs: number,
+  messages: WorkerCallErrorMessages
+): Promise<unknown> {
+  const requestId = generateWarmupId();
+  let message: string;
+  try {
+    message = JSON.stringify({
+      id: requestId,
+      protocol: PROTOCOL_ID,
+      method: 'call',
+      params: {
+        module: request.module,
+        functionName: request.functionName,
+        args: request.args,
+        kwargs: {},
+      },
+    });
+  } catch (error) {
+    throw new BridgeExecutionError(
+      `${messages.label} failed to encode request: ${error instanceof Error ? error.message : String(error)}`,
+      { cause: error instanceof Error ? error : undefined }
+    );
+  }
+
+  let response: string;
+  try {
+    response = await worker.transport.send(message, timeoutMs);
+  } catch (error) {
+    throw new BridgeExecutionError(
+      `${messages.label} failed to send: ${error instanceof Error ? error.message : String(error)}`,
+      { cause: error instanceof Error ? error : undefined }
+    );
+  }
+
+  let parsed: unknown;
+  try {
+    parsed = JSON.parse(response);
+  } catch (error) {
+    throw new BridgeExecutionError(`${messages.label} ${messages.invalidJsonMessage}`, {
+      cause: error instanceof Error ? error : undefined,
+    });
+  }
+
+  if (!parsed || typeof parsed !== 'object' || Array.isArray(parsed)) {
+    throw new BridgeExecutionError(
+      `${messages.label} ${messages.malformedEnvelopeMessage(requestId)}`
+    );
+  }
+
+  const envelope = parsed as { result?: unknown; error?: unknown };
+  if ('error' in envelope && envelope.error !== undefined && envelope.error !== null) {
+    const errorPayload = envelope.error;
+    if (errorPayload && typeof errorPayload === 'object' && !Array.isArray(errorPayload)) {
+      const err = errorPayload as { type?: unknown; message?: unknown };
+      const errType = typeof err.type === 'string' ? err.type : 'Error';
+      const errMessage = typeof err.message === 'string' ? err.message : 'Unknown error';
+      throw new BridgeExecutionError(
+        `${messages.label} ${messages.pythonErrorMessage(errType, errMessage)}`
+      );
+    }
+
+    throw new BridgeExecutionError(`${messages.label} ${messages.malformedErrorPayloadMessage}`);
+  }
+
+  if (!('result' in envelope)) {
+    throw new BridgeExecutionError(
+      `${messages.label} ${messages.malformedEnvelopeMessage(requestId)}`
+    );
+  }
+
+  return envelope.result;
+}
+
 // =============================================================================
 // ENVIRONMENT BUILDING
 // =============================================================================
@@ -513,13 +778,13 @@ function createWarmupCallback(
 function buildProcessEnv(options: ResolvedOptions): Record<string, string> {
   const allowedPrefixes = ['TYWRAP_'];
   const allowedKeys = new Set(['path', 'pythonpath', 'virtual_env', 'pythonhome']);
-  const env: Record<string, string> = {};
+  const env = createNullPrototypeEnv();
 
   // Copy allowed env vars from process.env
   if (options.inheritProcessEnv) {
     for (const [key, value] of Object.entries(process.env)) {
       if (value !== undefined) {
-        env[key] = value;
+        setEnvValue(env, key, value);
       }
     }
   } else {
@@ -528,15 +793,16 @@ function buildProcessEnv(options: ResolvedOptions): Record<string, string> {
         value !== undefined &&
         (allowedKeys.has(key.toLowerCase()) || allowedPrefixes.some(p => key.startsWith(p)))
       ) {
-        env[key] = value;
+        setEnvValue(env, key, value);
       }
     }
   }
 
   // Apply user overrides
   for (const [key, value] of Object.entries(options.env)) {
+    assertSafeEnvOverrideKey(key);
     if (value !== undefined) {
-      env[key] = value;
+      setEnvValue(env, key, value);
     }
   }
 
@@ -545,8 +811,8 @@ function buildProcessEnv(options: ResolvedOptions): Record<string, string> {
     const venv = resolveVirtualEnv(options.virtualEnv, options.cwd);
     env.VIRTUAL_ENV = venv.venvPath;
     const pathKey = getPathKey(env);
-    const currentPath = env[pathKey] ?? '';
-    env[pathKey] = `${venv.binDir}${delimiter}${currentPath}`;
+    const currentPath = getEnvValue(env, pathKey);
+    setPathValue(env, currentPath ? `${venv.binDir}${delimiter}${currentPath}` : venv.binDir);
   }
 
   // Add cwd to PYTHONPATH so Python can find modules in the working directory

package/src/runtime/optimized-node.ts

@@ -1,9 +1,9 @@
 /**
  * @deprecated Import from './node.js' instead.
  *
- * OptimizedNodeBridge has been unified with NodeBridge. The NodeBridge class
- * now supports both single-process mode (default) and multi-process pooling
- * via the minProcesses/maxProcesses options.
+ * Backward-compat shim for older deep imports. NodeBridge is the public API,
+ * and it now supports both single-process mode (default) and multi-process
+ * pooling via the minProcesses/maxProcesses options.
  *
  * Migration:
  * ```typescript
@@ -16,6 +16,10 @@
  * const bridge = new NodeBridge({ minProcesses: 2, maxProcesses: 4 });
  * ```
  *
- * This file is maintained for backward compatibility only.
+ * This file is not exposed through the package exports map and is maintained
+ * for backward compatibility only.
  */
-export { NodeBridge as OptimizedNodeBridge, type NodeBridgeOptions as ProcessPoolOptions } from './node.js';
+export {
+  NodeBridge as OptimizedNodeBridge,
+  type NodeBridgeOptions as ProcessPoolOptions,
+} from './node.js';

package/src/runtime/pooled-transport.ts

@@ -41,6 +41,11 @@ export interface PooledTransportOptions {
    * Use this for per-worker warmup (e.g., importing modules, running setup).
    */
   onWorkerReady?: (worker: PooledWorker) => Promise<void>;
+
+  /**
+   * Optional callback used only for background replacement workers.
+   */
+  onReplacementWorkerReady?: (worker: PooledWorker) => Promise<void>;
 }
 
 // =============================================================================
@@ -78,8 +83,12 @@ export interface PooledTransportOptions {
  * ```
  */
 export class PooledTransport extends BoundedContext implements Transport {
-  private readonly poolOptions: Omit<Required<PooledTransportOptions>, 'onWorkerReady'> & {
+  private readonly poolOptions: Omit<
+    Required<PooledTransportOptions>,
+    'onWorkerReady' | 'onReplacementWorkerReady'
+  > & {
     onWorkerReady?: (worker: PooledWorker) => Promise<void>;
+    onReplacementWorkerReady?: (worker: PooledWorker) => Promise<void>;
   };
   private pool?: WorkerPool;
 
@@ -102,6 +111,7 @@ export class PooledTransport extends BoundedContext implements Transport {
       queueTimeoutMs: options.queueTimeoutMs ?? 30000,
       maxConcurrentPerWorker: options.maxConcurrentPerWorker ?? 10,
       onWorkerReady: options.onWorkerReady,
+      onReplacementWorkerReady: options.onReplacementWorkerReady,
     };
   }
 
@@ -123,6 +133,7 @@ export class PooledTransport extends BoundedContext implements Transport {
       queueTimeoutMs: this.poolOptions.queueTimeoutMs,
       maxConcurrentPerWorker: this.poolOptions.maxConcurrentPerWorker,
       onWorkerReady: this.poolOptions.onWorkerReady,
+      onReplacementWorkerReady: this.poolOptions.onReplacementWorkerReady,
     });
 
     await this.pool.init();