makelore/electron/coding-runtime/pi/worker-process.ts

import {
  execFile,
  spawn,
  type ChildProcessWithoutNullStreams,
} from 'node:child_process';
import { platform } from 'node:os';
import { PiProcessError } from './process-errors';
import {
  PiRpcClient,
  type PiRpcCommand,
  type PiRpcEvent,
  type PiRpcRequestOptions,
  type PiRpcResponse,
} from './rpc-client';
import { StrictLfJsonlFramer } from './rpc-framer';

const DEFAULT_COMMAND_TIMEOUT_MS = 10_000;
const DEFAULT_SHUTDOWN_GRACE_MS = 3_000;
const DEFAULT_DIAGNOSTIC_BYTES = 16_000;
const ANSI_COLOR_PATTERN = new RegExp(`${String.fromCharCode(27)}\\[[0-9;]*m`, 'g');
const PI_INHERITED_ENV_KEYS = [
  'APPDATA',
  'COMSPEC',
  'HOME',
  'LANG',
  'LC_ALL',
  'LD_LIBRARY_PATH',
  'LOCALAPPDATA',
  'NODE_EXTRA_CA_CERTS',
  'PATH',
  'PATHEXT',
  'SSL_CERT_DIR',
  'SSL_CERT_FILE',
  'SYSTEMROOT',
  'TEMP',
  'TMP',
  'TMPDIR',
  'TZ',
  'USERPROFILE',
  'WINDIR',
] as const;

export type PiWorkerStopResult = {
  mode: 'not-started' | 'stdin-close' | 'forced-tree-kill';
  code: number | null;
  signal: NodeJS.Signals | null;
};

export type PiWorkerProcessOptions = {
  executablePath: string;
  cliPath: string;
  cwd: string;
  configDir: string;
  sessionDir: string;
  additionalArgs?: readonly string[];
  env?: NodeJS.ProcessEnv;
  sensitiveValues?: readonly string[];
  commandTimeoutMs?: number;
  shutdownGraceMs?: number;
  maxLineBytes?: number;
  diagnosticBytes?: number;
};

export function buildPiRpcArgs(
  sessionDir: string,
  additionalArgs: readonly string[] = [],
): string[] {
  return [
    '--mode', 'rpc',
    '--offline',
    '--session-dir', sessionDir,
    '--no-extensions',
    '--no-skills',
    '--no-prompt-templates',
    '--no-themes',
    '--no-context-files',
    '--no-approve',
    '--no-tools',
    ...additionalArgs,
  ];
}

export function sanitizePiDiagnostic(
  source: string,
  sensitiveValues: readonly string[] = [],
): string {
  let sanitized = source
    .replace(ANSI_COLOR_PATTERN, '')
    .replace(/(authorization\s*[:=]\s*(?:bearer\s+)?)[^\s,;]+/gi, '$1[REDACTED]')
    .replace(/((?:x-api-key|api[_-]?key|token|secret)\s*[:=]\s*)[^\s,;]+/gi, '$1[REDACTED]');
  const uniqueSensitiveValues = [...new Set(sensitiveValues.filter(Boolean))]
    .sort((left, right) => right.length - left.length);
  for (const value of uniqueSensitiveValues) {
    sanitized = sanitized.split(value).join('[REDACTED]');
  }
  return sanitized;
}

export function buildPiWorkerEnvironment(
  configDir: string,
  overlay: NodeJS.ProcessEnv = {},
  inherited: NodeJS.ProcessEnv = process.env,
): NodeJS.ProcessEnv {
  const env: NodeJS.ProcessEnv = {};
  for (const key of PI_INHERITED_ENV_KEYS) {
    const exact = inherited[key];
    if (exact !== undefined) {
      env[key] = exact;
      continue;
    }
    const matchingKey = Object.keys(inherited).find((candidate) => candidate.toUpperCase() === key);
    if (matchingKey && inherited[matchingKey] !== undefined) env[matchingKey] = inherited[matchingKey];
  }
  return {
    ...env,
    ...overlay,
    ELECTRON_RUN_AS_NODE: '1',
    PI_CODING_AGENT_DIR: configDir,
    PI_OFFLINE: '1',
    PI_TELEMETRY: '0',
  };
}

function assertSensitiveValuesAbsentFromArgs(
  args: readonly string[],
  sensitiveValues: readonly string[] = [],
): void {
  for (const value of sensitiveValues) {
    if (value && args.some((argument) => argument.includes(value))) {
      throw new Error('Pi worker arguments contain a sensitive value');
    }
  }
}

function boundedUtf8Tail(source: string, maxBytes: number): string {
  const bytes = Buffer.from(source, 'utf8');
  if (bytes.length <= maxBytes) return source;
  const decoder = new TextDecoder('utf-8', { fatal: true });
  for (let start = bytes.length - maxBytes; start < bytes.length; start += 1) {
    try {
      return decoder.decode(bytes.subarray(start));
    } catch {
      // A UTF-8 code point occupies at most four bytes; advance to its boundary.
    }
  }
  return '';
}

function positiveInteger(value: number | undefined, fallback: number, name: string): number {
  const resolved = value ?? fallback;
  if (!Number.isSafeInteger(resolved) || resolved <= 0) {
    throw new Error(`${name} must be a positive safe integer`);
  }
  return resolved;
}

function runExecutable(executable: string, args: readonly string[]): Promise<void> {
  return new Promise((resolve, reject) => {
    execFile(executable, [...args], { windowsHide: true }, (error) => {
      if (error) reject(error);
      else resolve();
    });
  });
}

export class PiWorkerProcess {
  private readonly options: PiWorkerProcessOptions;
  private readonly commandTimeoutMs: number;
  private readonly shutdownGraceMs: number;
  private readonly diagnosticBytes: number;
  private child: ChildProcessWithoutNullStreams | null = null;
  private rpc: PiRpcClient | null = null;
  private generationValue = 1;
  private diagnostic = '';
  private stdoutTail = '';
  private invalidation: PiProcessError | null = null;
  private stopping = false;
  private stopPromise: Promise<PiWorkerStopResult> | null = null;
  private exitResult: Promise<{ code: number | null; signal: NodeJS.Signals | null }> | null = null;

  constructor(options: PiWorkerProcessOptions) {
    this.options = options;
    this.commandTimeoutMs = positiveInteger(
      options.commandTimeoutMs,
      DEFAULT_COMMAND_TIMEOUT_MS,
      'commandTimeoutMs',
    );
    this.shutdownGraceMs = positiveInteger(
      options.shutdownGraceMs,
      DEFAULT_SHUTDOWN_GRACE_MS,
      'shutdownGraceMs',
    );
    this.diagnosticBytes = positiveInteger(
      options.diagnosticBytes,
      DEFAULT_DIAGNOSTIC_BYTES,
      'diagnosticBytes',
    );
  }

  get generation(): number {
    return this.generationValue;
  }

  get pendingCommandCount(): number {
    return this.rpc?.pendingCount ?? 0;
  }

  get stderrDiagnostic(): string {
    return this.diagnostic;
  }

  get protocolError(): PiProcessError | null {
    return this.invalidation?.code === 'PI_RPC_PROTOCOL_ERROR' ? this.invalidation : null;
  }

  async start(): Promise<this> {
    if (this.child) throw new Error('Pi worker process already started');
    const generation = this.generationValue;
    const args = buildPiRpcArgs(this.options.sessionDir, this.options.additionalArgs);
    assertSensitiveValuesAbsentFromArgs(args, this.options.sensitiveValues);
    const child = spawn(
      this.options.executablePath,
      [this.options.cliPath, ...args],
      {
        cwd: this.options.cwd,
        env: buildPiWorkerEnvironment(this.options.configDir, this.options.env),
        stdio: ['pipe', 'pipe', 'pipe'],
        windowsHide: true,
        detached: platform() !== 'win32',
      },
    );
    this.child = child;
    this.rpc = new PiRpcClient(child.stdin, {
      generation,
      defaultTimeoutMs: this.commandTimeoutMs,
      idPrefix: 'makelore-pi',
      onEventListenerError: (error) => {
        this.appendDiagnostic(`[event-listener] ${error instanceof Error ? error.message : String(error)}\n`);
      },
    });

    this.exitResult = new Promise((resolve) => {
      child.once('exit', (code, signal) => {
        resolve({ code, signal });
      });
    });
    child.once('close', (code, signal) => {
      if (!this.stopping && !this.invalidation) {
        this.invalidate(new PiProcessError(
          'PI_RPC_EXITED',
          `Pi worker exited unexpectedly (code=${code ?? 'null'}, signal=${signal ?? 'none'})`,
          { generation },
        ));
      }
    });

    const framer = new StrictLfJsonlFramer({
      maxLineBytes: this.options.maxLineBytes,
      onRecord: (record) => this.acceptRecord(record),
    });
    child.stdout.on('data', (chunk: Buffer) => {
      this.stdoutTail = boundedUtf8Tail(
        `${this.stdoutTail}${sanitizePiDiagnostic(chunk.toString('utf8'), this.options.sensitiveValues)}`,
        Math.min(this.diagnosticBytes, 2_048),
      );
      try {
        framer.push(chunk);
      } catch (error) {
        this.handleProtocolFailure(error);
      }
    });
    child.stdout.on('end', () => {
      try {
        framer.finish();
      } catch (error) {
        this.handleProtocolFailure(error);
      }
    });
    child.stderr.on('data', (chunk: Buffer) => this.appendDiagnostic(chunk.toString('utf8')));

    try {
      await new Promise<void>((resolve, reject) => {
        child.once('spawn', resolve);
        child.once('error', reject);
      });
    } catch (error) {
      const failure = new PiProcessError('PI_WORKER_START_FAILED', 'Could not start Pi worker', {
        cause: error,
        generation,
        diagnostic: this.diagnostic,
      });
      this.invalidate(failure);
      throw failure;
    }
    return this;
  }

  request<T = unknown>(
    command: PiRpcCommand,
    options?: PiRpcRequestOptions,
  ): Promise<PiRpcResponse<T>> {
    if (!this.rpc) {
      return Promise.reject(new PiProcessError(
        'PI_WORKER_START_FAILED',
        'Pi worker has not started',
        { generation: this.generationValue },
      ));
    }
    return this.rpc.request<T>(command, options);
  }

  subscribe(listener: (event: PiRpcEvent) => void): () => void {
    if (!this.rpc) throw new Error('Pi worker has not started');
    return this.rpc.subscribe(listener);
  }

  stop(): Promise<PiWorkerStopResult> {
    if (!this.stopPromise) this.stopPromise = this.performStop();
    return this.stopPromise;
  }

  private acceptRecord(record: unknown): void {
    try {
      this.rpc?.accept(record);
    } catch (error) {
      this.handleProtocolFailure(error);
    }
  }

  private handleProtocolFailure(error: unknown): void {
    if (this.invalidation) return;
    const message = error instanceof Error ? error.message : 'Pi RPC protocol failure';
    this.appendDiagnostic(`[stdout-protocol] ${this.stdoutTail}\n`);
    const failure = new PiProcessError('PI_RPC_PROTOCOL_ERROR', message, {
      cause: error,
      generation: this.generationValue,
      diagnostic: this.diagnostic,
    });
    this.invalidate(failure);
    void this.forceKillTree();
  }

  private invalidate(error: PiProcessError): void {
    if (this.invalidation) return;
    this.invalidation = error;
    this.rpc?.invalidate(error);
    this.generationValue += 1;
  }

  private appendDiagnostic(source: string): void {
    const sanitized = sanitizePiDiagnostic(source, this.options.sensitiveValues);
    this.diagnostic = boundedUtf8Tail(`${this.diagnostic}${sanitized}`, this.diagnosticBytes);
  }

  private async performStop(): Promise<PiWorkerStopResult> {
    const child = this.child;
    const exitResult = this.exitResult;
    if (!child || !exitResult) {
      return { mode: 'not-started', code: null, signal: null };
    }

    this.stopping = true;
    this.invalidate(new PiProcessError('PI_WORKER_STOPPED', 'Pi worker stopped', {
      generation: this.generationValue,
    }));
    child.stdin.end();
    const graceful = await this.waitForExit(this.shutdownGraceMs);
    if (graceful) return { mode: 'stdin-close', ...graceful };

    await this.forceKillTree();
    const forced = await this.waitForExit(this.shutdownGraceMs);
    if (forced) return { mode: 'forced-tree-kill', ...forced };
    throw new PiProcessError('PI_WORKER_STOP_FAILED', 'Pi worker did not exit after forced tree kill', {
      generation: this.generationValue,
      diagnostic: this.diagnostic,
    });
  }

  private async waitForExit(
    timeoutMs: number,
  ): Promise<{ code: number | null; signal: NodeJS.Signals | null } | null> {
    if (!this.exitResult) return null;
    let timer: ReturnType<typeof setTimeout> | undefined;
    try {
      return await Promise.race([
        this.exitResult,
        new Promise<null>((resolve) => {
          timer = setTimeout(() => resolve(null), timeoutMs);
        }),
      ]);
    } finally {
      if (timer) clearTimeout(timer);
    }
  }

  private async forceKillTree(): Promise<void> {
    const child = this.child;
    if (!child?.pid) return;
    if (platform() === 'win32') {
      await runExecutable('taskkill.exe', ['/pid', String(child.pid), '/t', '/f']).catch(() => undefined);
      return;
    }
    try {
      process.kill(-child.pid, 'SIGKILL');
    } catch {
      child.kill('SIGKILL');
    }
  }
}