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

import {
  execFile,
  spawn,
  type ChildProcessWithoutNullStreams,
} from 'node:child_process';
import { realpathSync } from 'node:fs';
import { platform } from 'node:os';
import { logger } from '../../utils/logger';
import type { CodingRuntimeDisposeReason } from '../contracts';
import { PiProcessError, type PiProcessErrorCode } 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 PiWorkerProofFailure = 'unexpected_exit' | 'protocol_invalidation';

export type PiWorkerStopReason = CodingRuntimeDisposeReason
  | 'app_shutdown'
  | 'idle_eviction'
  | 'queued_suspension'
  | 'stale_resource_rebuild'
  | 'recover'
  | 'process_capacity_reopen'
  | 'protocol_invalidation'
  | 'unexpected_exit_cleanup'
  | 'open_failure'
  | 'session_binding_mismatch'
  | 'subagent_complete'
  | 'subagent_abort'
  | 'subagent_open_failure'
  | 'test_injection';

const PI_WORKER_STOP_REASONS = new Set<PiWorkerStopReason>([
  'app_shutdown',
  'idle_eviction',
  'queued_suspension',
  'stale_resource_rebuild',
  'recover',
  'background_sleep',
  'project_deactivated',
  'project_removed',
  'conversation_deleted',
  'auth_cleanup',
  'model_reconfiguration',
  'process_capacity_reopen',
  'fork_replacement',
  'protocol_invalidation',
  'unexpected_exit_cleanup',
  'open_failure',
  'session_binding_mismatch',
  'subagent_complete',
  'subagent_abort',
  'subagent_open_failure',
  'test_injection',
]);

export type PiWorkerLifecycleEvent = {
  classification: 'unexpected_exit' | 'protocol_invalidation' | 'intentional_stop';
  stage: 'protocol' | 'close' | 'stop_requested' | 'stop_completed' | 'stop_failed';
  conversationId?: string;
  generation: number;
  code: PiProcessErrorCode;
  reason?: PiWorkerStopReason;
  exitCode: number | null;
  signal: NodeJS.Signals | null;
  diagnostic?: string;
};

export type PiWorkerProcessOptions = {
  executablePath: string;
  cliPath: string;
  cwd: string;
  configDir: string;
  sessionDir: string;
  tools?: readonly string[];
  additionalArgs?: readonly string[];
  env?: NodeJS.ProcessEnv;
  sensitiveValues?: readonly string[];
  commandTimeoutMs?: number;
  shutdownGraceMs?: number;
  maxLineBytes?: number;
  diagnosticBytes?: number;
  conversationId?: string;
  workerGeneration?: number;
  onLifecycleEvent?(event: PiWorkerLifecycleEvent): void;
};

/** Tools available to every managed parent worker before plugin materialization. */
export const PI_CORE_TOOL_NAMES = Object.freeze([
  'read', 'bash', 'edit', 'write', 'grep', 'find', 'ls', 'ask_user', 'subagent',
  'agent_browser',
  'task_state', 'changed_file', 'runtime_context',
] as const);

/** Compatibility alias used by the shared Agent Server transport. */
export const DEFAULT_PI_RPC_TOOLS = PI_CORE_TOOL_NAMES;

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

function diagnosticPathAliases(value: string | undefined): string[] {
  if (!value) return [];
  try {
    const canonical = realpathSync.native(value);
    return canonical === value ? [value] : [value, canonical];
  } catch {
    return [value];
  }
}

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 readonly diagnosticSensitiveValues: string[];
  private child: ChildProcessWithoutNullStreams | null = null;
  private rpc: PiRpcClient | null = null;
  private generationValue: number;
  private diagnostic = '';
  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;
  private readonly invalidationListeners = new Set<(error: PiProcessError) => void>();

  constructor(options: PiWorkerProcessOptions) {
    this.options = options;
    this.generationValue = positiveInteger(options.workerGeneration, 1, 'workerGeneration');
    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',
    );
    const diagnosticPaths = [
      options.cwd,
      options.cliPath,
      options.configDir,
      options.sessionDir,
      options.env?.HOME,
      options.env?.USERPROFILE,
      options.env?.APPDATA,
      options.env?.LOCALAPPDATA,
    ].flatMap(diagnosticPathAliases);
    this.diagnosticSensitiveValues = [...new Set([
      ...(options.sensitiveValues ?? []),
      ...diagnosticPaths,
    ].filter((value): value is string => Boolean(value)))];
  }

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

  get processId(): number | undefined {
    return this.child?.pid;
  }

  get isRunning(): boolean {
    return this.child !== null
      && this.child.exitCode === null
      && this.child.signalCode === null;
  }

  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,
      this.options.tools,
    );
    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) {
        const failure = new PiProcessError(
          'PI_RPC_EXITED',
          `Pi worker exited unexpectedly (code=${code ?? 'null'}, signal=${signal ?? 'none'})`,
          {
            generation,
            diagnostic: this.diagnostic || undefined,
            exitCode: code,
            signal,
          },
        );
        this.recordLifecycle({
          classification: 'unexpected_exit',
          stage: 'close',
          generation,
          code: failure.code,
          exitCode: code,
          signal,
          ...(this.diagnostic ? { diagnostic: this.diagnostic } : {}),
        });
        this.invalidate(failure);
      }
    });

    const framer = new StrictLfJsonlFramer({
      maxLineBytes: this.options.maxLineBytes,
      onRecord: (record) => this.acceptRecord(record),
    });
    child.stdout.on('data', (chunk: Buffer) => {
      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);
  }

  delayNextResponseForProof(commandType: string, delayMs: number): void {
    if (!this.rpc) throw new Error('Pi worker has not started');
    this.rpc.delayNextResponseForProof(commandType, delayMs);
  }

  send(command: PiRpcCommand): Promise<void> {
    if (!this.rpc) {
      return Promise.reject(new PiProcessError(
        'PI_WORKER_START_FAILED',
        'Pi worker has not started',
        { generation: this.generationValue },
      ));
    }
    return this.rpc.send(command);
  }

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

  subscribeInvalidation(listener: (error: PiProcessError) => void): () => void {
    this.invalidationListeners.add(listener);
    return () => this.invalidationListeners.delete(listener);
  }

  stop(reason: PiWorkerStopReason): Promise<PiWorkerStopResult> {
    if (!PI_WORKER_STOP_REASONS.has(reason)) {
      return Promise.reject(new Error('Pi worker stop reason is required'));
    }
    if (!this.stopPromise) this.stopPromise = this.performStop(reason);
    return this.stopPromise;
  }

  async injectFailureForProof(failure: PiWorkerProofFailure): Promise<void> {
    if (!this.isRunning) throw new Error('Pi worker proof failure requires a running process');
    this.appendDiagnostic(
      '[release-proof] worker failure Authorization: Bearer packaged-proof-secret\n',
    );
    if (failure === 'protocol_invalidation') {
      this.handleProtocolFailure(new Error('Injected strict JSONL protocol failure'));
      return;
    }
    await this.forceKillTree();
  }

  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] ${message}\n`);
    const failure = new PiProcessError('PI_RPC_PROTOCOL_ERROR', message, {
      cause: error,
      generation: this.generationValue,
      diagnostic: this.diagnostic,
    });
    this.recordLifecycle({
      classification: 'protocol_invalidation',
      stage: 'protocol',
      generation: this.generationValue,
      code: failure.code,
      exitCode: this.child?.exitCode ?? null,
      signal: this.child?.signalCode ?? null,
      ...(this.diagnostic ? { 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;
    for (const listener of this.invalidationListeners) {
      try {
        listener(error);
      } catch (listenerError) {
        this.appendDiagnostic(
          `[invalidation-listener] ${listenerError instanceof Error ? listenerError.message : String(listenerError)}\n`,
        );
      }
    }
  }

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

  private async performStop(reason: PiWorkerStopReason): Promise<PiWorkerStopResult> {
    const child = this.child;
    const exitResult = this.exitResult;
    const generation = this.invalidation?.generation ?? this.generationValue;
    const classification = this.invalidation?.code === 'PI_RPC_PROTOCOL_ERROR'
      ? 'protocol_invalidation'
      : this.invalidation?.code === 'PI_RPC_EXITED'
        ? 'unexpected_exit'
        : 'intentional_stop';
    this.recordLifecycle({
      classification,
      stage: 'stop_requested',
      generation,
      code: this.invalidation?.code ?? 'PI_WORKER_STOPPED',
      reason,
      exitCode: child?.exitCode ?? null,
      signal: child?.signalCode ?? null,
      ...(this.diagnostic ? { diagnostic: this.diagnostic } : {}),
    });
    if (!child || !exitResult) {
      this.recordLifecycle({
        classification,
        stage: 'stop_completed',
        generation,
        code: this.invalidation?.code ?? 'PI_WORKER_STOPPED',
        reason,
        exitCode: null,
        signal: null,
        ...(this.diagnostic ? { diagnostic: this.diagnostic } : {}),
      });
      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) {
      this.recordLifecycle({
        classification,
        stage: 'stop_completed',
        generation,
        code: this.invalidation?.code ?? 'PI_WORKER_STOPPED',
        reason,
        exitCode: graceful.code,
        signal: graceful.signal,
        ...(this.diagnostic ? { diagnostic: this.diagnostic } : {}),
      });
      return { mode: 'stdin-close', ...graceful };
    }

    await this.forceKillTree();
    const forced = await this.waitForExit(this.shutdownGraceMs);
    if (forced) {
      this.recordLifecycle({
        classification,
        stage: 'stop_completed',
        generation,
        code: this.invalidation?.code ?? 'PI_WORKER_STOPPED',
        reason,
        exitCode: forced.code,
        signal: forced.signal,
        ...(this.diagnostic ? { diagnostic: this.diagnostic } : {}),
      });
      return { mode: 'forced-tree-kill', ...forced };
    }
    this.recordLifecycle({
      classification,
      stage: 'stop_failed',
      generation,
      code: 'PI_WORKER_STOP_FAILED',
      reason,
      exitCode: child.exitCode,
      signal: child.signalCode,
      ...(this.diagnostic ? { diagnostic: this.diagnostic } : {}),
    });
    throw new PiProcessError('PI_WORKER_STOP_FAILED', 'Pi worker did not exit after forced tree kill', {
      generation: this.generationValue,
      diagnostic: this.diagnostic,
    });
  }

  private recordLifecycle(event: PiWorkerLifecycleEvent): void {
    const diagnostic = event.diagnostic
      ? boundedUtf8Tail(
          sanitizePiDiagnostic(event.diagnostic, this.diagnosticSensitiveValues),
          this.diagnosticBytes,
        )
      : undefined;
    const safeEvent: PiWorkerLifecycleEvent = {
      ...event,
      ...(this.options.conversationId ? { conversationId: this.options.conversationId } : {}),
      ...(diagnostic ? { diagnostic } : {}),
    };
    try {
      this.options.onLifecycleEvent?.(structuredClone(safeEvent));
    } catch {
      // Lifecycle observers must not affect process cleanup.
    }
    logger.warn('[PiWorkerLifecycle]', safeEvent);
  }

  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');
    }
  }
}