makelore/electron/coding-runtime/pi/extension-ui-projector.ts

import type { PiRpcEvent } from './rpc-client';

const MAX_MESSAGE_LENGTH = 2_000;
const MAX_TITLE_LENGTH = 256;
const MAX_EDITOR_TEXT_LENGTH = 64 * 1024;
const MAX_WIDGET_LINES = 32;
const MAX_WIDGET_LINE_LENGTH = 512;
const MAX_DIAGNOSTICS = 256;

export type PiExtensionUiProjection =
  | { kind: 'notify'; conversationId: string; message: string; level: 'info' | 'warning' | 'error' }
  | { kind: 'status'; conversationId: string; key: string; text?: string }
  | {
    kind: 'widget';
    conversationId: string;
    key: string;
    lines?: string[];
    placement?: 'aboveEditor' | 'belowEditor';
  }
  | { kind: 'title'; conversationId: string; title: string }
  | { kind: 'editor-text'; conversationId: string; text: string; draftRevision: number };

export interface PiExtensionUiDiagnostic {
  method: string;
  reason: 'unsupported-ui-method' | 'invalid-ui-payload' | 'stale-draft-revision';
}

export interface PiExtensionUiProjectorOptions {
  getDraftRevision(conversationId: string): number;
  knownStatusKeys?: readonly string[];
  knownWidgetKeys?: readonly string[];
}

interface RunProjectionState {
  runId: string;
  generation: number;
  draftRevision: number;
}

function bounded(value: string, length: number): string {
  return value.length <= length ? value : `${value.slice(0, length - 1)}…`;
}

export class PiExtensionUiProjector {
  private readonly getDraftRevision: (conversationId: string) => number;
  private readonly knownStatusKeys: Set<string>;
  private readonly knownWidgetKeys: Set<string>;
  private readonly runs = new Map<string, RunProjectionState>();
  private readonly diagnostics: PiExtensionUiDiagnostic[] = [];

  constructor(options: PiExtensionUiProjectorOptions) {
    this.getDraftRevision = options.getDraftRevision;
    this.knownStatusKeys = new Set(options.knownStatusKeys ?? ['makelore.write-lease']);
    this.knownWidgetKeys = new Set(options.knownWidgetKeys ?? []);
  }

  beginRun(conversationId: string, generation: number, runId: string): void {
    this.runs.set(conversationId, {
      generation,
      runId,
      draftRevision: this.getDraftRevision(conversationId),
    });
  }

  replaceGeneration(conversationId: string, generation: number, runId: string): void {
    const current = this.runs.get(conversationId);
    if (current?.runId === runId) current.generation = generation;
  }

  endRun(conversationId: string, runId: string): void {
    if (this.runs.get(conversationId)?.runId === runId) this.runs.delete(conversationId);
  }

  getDiagnostics(): PiExtensionUiDiagnostic[] {
    return structuredClone(this.diagnostics);
  }

  project(
    conversationId: string,
    generation: number,
    runId: string,
    event: PiRpcEvent,
  ): PiExtensionUiProjection | null {
    if (event.type !== 'extension_ui_request' || typeof event.method !== 'string') return null;
    const run = this.runs.get(conversationId);
    if (!run || run.runId !== runId || run.generation !== generation) return null;
    if (event.method === 'notify'
      && typeof event.message === 'string'
      && ['info', 'warning', 'error', undefined].includes(event.notifyType as string | undefined)) {
      return {
        kind: 'notify',
        conversationId,
        message: bounded(event.message, MAX_MESSAGE_LENGTH),
        level: event.notifyType === 'warning' || event.notifyType === 'error'
          ? event.notifyType
          : 'info',
      };
    }
    if (event.method === 'setStatus'
      && typeof event.statusKey === 'string'
      && (event.statusText === undefined || typeof event.statusText === 'string')) {
      if (!this.knownStatusKeys.has(event.statusKey)) return this.unsupported(event.method);
      return {
        kind: 'status',
        conversationId,
        key: event.statusKey,
        ...(typeof event.statusText === 'string'
          ? { text: bounded(event.statusText, MAX_MESSAGE_LENGTH) }
          : {}),
      };
    }
    if (event.method === 'setWidget'
      && typeof event.widgetKey === 'string'
      && (event.widgetLines === undefined || Array.isArray(event.widgetLines))) {
      if (!this.knownWidgetKeys.has(event.widgetKey)) return this.unsupported(event.method);
      if (Array.isArray(event.widgetLines)
        && !event.widgetLines.every((line) => typeof line === 'string')) {
        return this.invalid(event.method);
      }
      return {
        kind: 'widget',
        conversationId,
        key: event.widgetKey,
        ...(Array.isArray(event.widgetLines)
          ? {
            lines: event.widgetLines.slice(0, MAX_WIDGET_LINES)
              .map((line) => bounded(line as string, MAX_WIDGET_LINE_LENGTH)),
          }
          : {}),
        ...(event.widgetPlacement === 'aboveEditor' || event.widgetPlacement === 'belowEditor'
          ? { placement: event.widgetPlacement }
          : {}),
      };
    }
    if (event.method === 'setTitle' && typeof event.title === 'string') {
      return { kind: 'title', conversationId, title: bounded(event.title, MAX_TITLE_LENGTH) };
    }
    if (event.method === 'set_editor_text' && typeof event.text === 'string') {
      if (this.getDraftRevision(conversationId) !== run.draftRevision) {
        this.recordDiagnostic({ method: event.method, reason: 'stale-draft-revision' });
        return null;
      }
      return {
        kind: 'editor-text',
        conversationId,
        text: bounded(event.text, MAX_EDITOR_TEXT_LENGTH),
        draftRevision: run.draftRevision,
      };
    }
    if (['select', 'confirm', 'input', 'editor'].includes(event.method)) return null;
    return this.invalid(event.method);
  }

  private unsupported(method: string): null {
    this.recordDiagnostic({ method: bounded(method, 128), reason: 'unsupported-ui-method' });
    return null;
  }

  private invalid(method: string): null {
    this.recordDiagnostic({ method: bounded(method, 128), reason: 'invalid-ui-payload' });
    return null;
  }

  private recordDiagnostic(diagnostic: PiExtensionUiDiagnostic): void {
    this.diagnostics.push(diagnostic);
    if (this.diagnostics.length > MAX_DIAGNOSTICS) this.diagnostics.shift();
  }
}