makelore/electron/coding-runtime/pi/rpc-client.ts

import type { Writable } from 'node:stream';
import { PiProcessError } from './process-errors';

export type PiRpcCommand = {
  type: string;
  [key: string]: unknown;
};

export type PiRpcResponse<T = unknown> = {
  type: 'response';
  id: string;
  command?: string;
  success: boolean;
  data?: T;
  error?: string;
};

export type PiRpcEvent = Record<string, unknown> & { type: string };
export type PiRpcRetryPolicy = 'none' | 'read-only-once';

export type PiRpcRequestOptions = {
  signal?: AbortSignal;
  timeoutMs?: number;
  retry?: PiRpcRetryPolicy;
};

type PendingRequest = {
  commandType: string;
  resolve(response: PiRpcResponse): void;
  reject(error: PiProcessError): void;
  cancel(): void;
};

type PiRpcClientOptions = {
  generation: number;
  defaultTimeoutMs?: number;
  idPrefix?: string;
  onEventListenerError?(error: unknown): void;
};

const READ_ONLY_COMMANDS = new Set([
  'get_available_models',
  'get_available_thinking_levels',
  'get_commands',
  'get_entries',
  'get_fork_messages',
  'get_last_assistant_text',
  'get_messages',
  'get_session_stats',
  'get_state',
  'get_tree',
]);

function recordValue(value: unknown): value is Record<string, unknown> {
  return Boolean(value) && typeof value === 'object' && !Array.isArray(value);
}

function responseValue(value: Record<string, unknown>): value is PiRpcResponse {
  return value.type === 'response'
    && typeof value.id === 'string'
    && typeof value.success === 'boolean';
}

export class PiRpcClient {
  private readonly writable: Writable;
  private readonly generation: number;
  private readonly defaultTimeoutMs: number;
  private readonly idPrefix: string;
  private readonly onEventListenerError: ((error: unknown) => void) | undefined;
  private readonly pending = new Map<string, PendingRequest>();
  private readonly retiredIds = new Set<string>();
  private readonly listeners = new Set<(event: PiRpcEvent) => void>();
  private sequence = 0;
  private invalidated: PiProcessError | null = null;

  constructor(writable: Writable, options: PiRpcClientOptions) {
    if (!Number.isSafeInteger(options.generation) || options.generation <= 0) {
      throw new Error('generation must be a positive safe integer');
    }
    const defaultTimeoutMs = options.defaultTimeoutMs ?? 10_000;
    if (!Number.isSafeInteger(defaultTimeoutMs) || defaultTimeoutMs <= 0) {
      throw new Error('defaultTimeoutMs must be a positive safe integer');
    }
    this.writable = writable;
    this.generation = options.generation;
    this.defaultTimeoutMs = defaultTimeoutMs;
    this.idPrefix = options.idPrefix ?? 'makelore';
    this.onEventListenerError = options.onEventListenerError;
    this.writable.on('error', (error) => {
      this.invalidate(new PiProcessError('PI_RPC_WRITE_FAILED', 'Pi RPC stdin failed', {
        cause: error,
        generation: this.generation,
      }));
    });
  }

  get pendingCount(): number {
    return this.pending.size;
  }

  subscribe(listener: (event: PiRpcEvent) => void): () => void {
    this.listeners.add(listener);
    return () => this.listeners.delete(listener);
  }

  async request<T = unknown>(
    command: PiRpcCommand,
    options: PiRpcRequestOptions = {},
  ): Promise<PiRpcResponse<T>> {
    if (!command.type) throw new Error('Pi RPC command type is required');
    const retry = options.retry ?? 'none';
    if (retry === 'read-only-once' && !READ_ONLY_COMMANDS.has(command.type)) {
      throw new Error(`Pi RPC ${command.type} is not a retryable read-only command`);
    }

    const attempts = retry === 'read-only-once' ? 2 : 1;
    let lastError: unknown;
    for (let attempt = 0; attempt < attempts; attempt += 1) {
      try {
        return await this.requestOnce<T>(command, options);
      } catch (error) {
        lastError = error;
        if (!(error instanceof PiProcessError) || error.code !== 'PI_RPC_TIMEOUT') throw error;
        if (this.invalidated || options.signal?.aborted) throw error;
      }
    }
    throw lastError;
  }

  async send(command: PiRpcCommand): Promise<void> {
    if (!command.type) throw new Error('Pi RPC command type is required');
    if (this.invalidated) throw this.invalidated;
    try {
      await this.write(`${JSON.stringify(command)}\n`);
    } catch (error) {
      throw new PiProcessError('PI_RPC_WRITE_FAILED', `Could not write Pi RPC ${command.type}`, {
        cause: error,
        generation: this.generation,
      });
    }
  }

  accept(value: unknown): void {
    if (!recordValue(value) || typeof value.type !== 'string') {
      throw new PiProcessError('PI_RPC_PROTOCOL_ERROR', 'Pi RPC record must be an object with a type');
    }
    if (value.type === 'response') {
      if (!responseValue(value)) {
        throw new PiProcessError('PI_RPC_PROTOCOL_ERROR', 'Pi RPC response shape is invalid');
      }
      if (this.retiredIds.delete(value.id)) return;
      const pending = this.pending.get(value.id);
      if (!pending) {
        throw new PiProcessError('PI_RPC_PROTOCOL_ERROR', `Pi RPC response used unknown id ${value.id}`);
      }
      pending.cancel();
      this.pending.delete(value.id);
      if (value.success) pending.resolve(value);
      else {
        pending.reject(new PiProcessError(
          'PI_RPC_RESPONSE_ERROR',
          `Pi RPC ${value.command ?? pending.commandType} failed: ${value.error ?? 'unknown error'}`,
          { generation: this.generation },
        ));
      }
      return;
    }

    for (const listener of this.listeners) {
      try {
        listener(value as PiRpcEvent);
      } catch (error) {
        try {
          this.onEventListenerError?.(error);
        } catch {
          // Diagnostics must not affect transport state.
        }
      }
    }
  }

  invalidate(error: PiProcessError): void {
    if (this.invalidated) return;
    this.invalidated = error;
    for (const [id, pending] of this.pending) {
      pending.cancel();
      this.retire(id);
      pending.reject(error);
    }
    this.pending.clear();
  }

  private async requestOnce<T>(
    command: PiRpcCommand,
    options: PiRpcRequestOptions,
  ): Promise<PiRpcResponse<T>> {
    if (this.invalidated) throw this.invalidated;
    if (options.signal?.aborted) {
      throw new PiProcessError('PI_RPC_ABORTED', `Pi RPC ${command.type} was aborted`, {
        generation: this.generation,
      });
    }

    const timeoutMs = options.timeoutMs ?? this.defaultTimeoutMs;
    if (!Number.isSafeInteger(timeoutMs) || timeoutMs <= 0) {
      throw new Error('timeoutMs must be a positive safe integer');
    }
    const id = `${this.idPrefix}-${this.generation}-${++this.sequence}`;

    const response = new Promise<PiRpcResponse<T>>((resolve, reject) => {
      const timeout = setTimeout(() => {
        this.pending.delete(id);
        this.retire(id);
        reject(new PiProcessError(
          'PI_RPC_TIMEOUT',
          `Pi RPC ${command.type} timed out after ${timeoutMs}ms`,
          { generation: this.generation },
        ));
      }, timeoutMs);
      const abort = (): void => {
        this.pending.delete(id);
        this.retire(id);
        clearTimeout(timeout);
        reject(new PiProcessError('PI_RPC_ABORTED', `Pi RPC ${command.type} was aborted`, {
          generation: this.generation,
        }));
      };
      options.signal?.addEventListener('abort', abort, { once: true });
      this.pending.set(id, {
        commandType: command.type,
        resolve: (value) => resolve(value as PiRpcResponse<T>),
        reject,
        cancel: () => {
          clearTimeout(timeout);
          options.signal?.removeEventListener('abort', abort);
        },
      });
    });

    const write = this.write(`${JSON.stringify({ ...command, id })}\n`).catch((error) => {
      throw new PiProcessError('PI_RPC_WRITE_FAILED', `Could not write Pi RPC ${command.type}`, {
        cause: error,
        generation: this.generation,
      });
    });
    try {
      const [, record] = await Promise.all([write, response]);
      return record;
    } catch (error) {
      const pending = this.pending.get(id);
      if (pending) {
        pending.cancel();
        this.pending.delete(id);
        this.retire(id);
        pending.reject(error instanceof PiProcessError
          ? error
          : new PiProcessError('PI_RPC_WRITE_FAILED', `Could not write Pi RPC ${command.type}`, {
            cause: error,
            generation: this.generation,
          }));
      }
      throw error;
    }
  }

  private write(line: string): Promise<void> {
    return new Promise((resolve, reject) => {
      this.writable.write(line, (error?: Error | null) => {
        if (error) reject(error);
        else resolve();
      });
    });
  }

  private retire(id: string): void {
    this.retiredIds.add(id);
    if (this.retiredIds.size <= 256) return;
    const oldest = this.retiredIds.values().next().value as string | undefined;
    if (oldest) this.retiredIds.delete(oldest);
  }
}