makelore/electron/coding-runtime/pi/interaction.ts

import type {
  ConversationInteraction,
  ConversationInteractionResponse,
} from '../contracts';
import type { PiRpcCommand, PiRpcEvent } from './rpc-client';
import type { PiGenerationResourceInput, PiWorkerPoolState } from './worker-pool';

interface PiInteractionTransport {
  getState(conversationId: string): PiWorkerPoolState | null;
  getActiveRun(conversationId: string): { runId: string; generation: number } | null;
  send(conversationId: string, command: PiRpcCommand): Promise<void>;
  trackGenerationResource(input: PiGenerationResourceInput): () => void;
}

interface StoredInteraction {
  interaction: ConversationInteraction;
  generation: number;
  labels: Map<string, string>;
  phase: 'pending' | 'responding';
  untrack(): void;
}

function dialogEvent(event: PiRpcEvent): event is PiRpcEvent & {
  id: string;
  method: 'select' | 'confirm' | 'input' | 'editor';
  title: string;
} {
  return event.type === 'extension_ui_request'
    && typeof event.id === 'string'
    && ['select', 'confirm', 'input', 'editor'].includes(String(event.method))
    && typeof event.title === 'string';
}

export class PiInteractionStore {
  private readonly pending = new Map<string, StoredInteraction>();

  constructor(
    private readonly transport: PiInteractionTransport,
    private readonly onChange: (interaction: ConversationInteraction) => void,
  ) {}

  list(conversationId?: string): ConversationInteraction[] {
    return [...this.pending.values()]
      .map(({ interaction }) => interaction)
      .filter((interaction) => !conversationId || interaction.conversationId === conversationId)
      .map((interaction) => structuredClone(interaction));
  }

  open(
    conversationId: string,
    generation: number,
    runId: string,
    event: PiRpcEvent,
  ): ConversationInteraction | null {
    if (!dialogEvent(event)) return null;
    const state = this.transport.getState(conversationId);
    const active = this.transport.getActiveRun(conversationId);
    if (state?.generation !== generation
      || active?.generation !== generation
      || active.runId !== runId) return null;
    const key = this.key(conversationId, event.id);
    if (this.pending.has(key)) throw new Error(`Duplicate Pi interaction id: ${event.id}`);
    const labels = new Map<string, string>();
    const options = event.method === 'select' && Array.isArray(event.options)
      ? event.options.flatMap((label, index) => {
        if (typeof label !== 'string') return [];
        const id = `${event.id}:option:${index}`;
        labels.set(id, label);
        return [{ id, label }];
      })
      : undefined;
    const interaction: ConversationInteraction = {
      id: event.id,
      conversationId,
      runId,
      kind: event.method,
      title: event.title,
      ...(typeof event.message === 'string' ? { message: event.message } : {}),
      ...(options ? { options } : {}),
      status: 'pending',
    };
    const stored: StoredInteraction = {
      interaction,
      generation,
      labels,
      phase: 'pending',
      untrack: () => undefined,
    };
    stored.untrack = this.transport.trackGenerationResource({
      conversationId,
      kind: 'interaction',
      id: event.id,
      cancel: () => {
        void this.cancelStored(stored, false);
      },
    });
    this.pending.set(key, stored);
    return structuredClone(interaction);
  }

  async respond(conversationId: string, response: ConversationInteractionResponse): Promise<ConversationInteraction> {
    const stored = this.pending.get(this.key(conversationId, response.interactionId));
    if (!stored) throw new Error('Pi interaction is not pending');
    if (stored.phase === 'responding') throw new Error('Pi interaction response is already in progress');
    const state = this.transport.getState(conversationId);
    const active = this.transport.getActiveRun(conversationId);
    if (state?.generation !== stored.generation
      || active?.generation !== stored.generation
      || active.runId !== stored.interaction.runId) {
      await this.cancelStored(stored, false);
      throw new Error('Pi interaction belongs to a stale worker run');
    }

    let command: PiRpcCommand;
    if ('cancelled' in response) {
      command = { type: 'extension_ui_response', id: stored.interaction.id, cancelled: true };
    } else if (stored.interaction.kind === 'select' && 'optionId' in response) {
      const value = stored.labels.get(response.optionId);
      if (value === undefined) throw new Error('Pi interaction option is invalid');
      command = { type: 'extension_ui_response', id: stored.interaction.id, value };
    } else if (stored.interaction.kind === 'confirm' && 'confirmed' in response) {
      command = { type: 'extension_ui_response', id: stored.interaction.id, confirmed: response.confirmed };
    } else if ((stored.interaction.kind === 'input' || stored.interaction.kind === 'editor')
      && 'value' in response) {
      command = { type: 'extension_ui_response', id: stored.interaction.id, value: response.value };
    } else {
      throw new Error('Pi interaction response does not match its kind');
    }
    stored.phase = 'responding';
    try {
      await this.transport.send(conversationId, command);
    } catch (error) {
      if (this.pending.get(this.key(conversationId, stored.interaction.id)) === stored) {
        stored.phase = 'pending';
      }
      throw error;
    }
    if (this.pending.get(this.key(conversationId, stored.interaction.id)) !== stored) {
      throw new Error('Pi interaction belongs to a stale worker run');
    }
    return this.finish(stored, 'cancelled' in response
      ? 'cancelled'
      : stored.interaction.kind === 'confirm' && 'confirmed' in response && !response.confirmed
        ? 'rejected'
        : 'answered');
  }

  async cancelRun(conversationId: string, runId: string, notifyWorker: boolean): Promise<void> {
    const targets = [...this.pending.values()].filter(({ interaction }) => (
      interaction.conversationId === conversationId && interaction.runId === runId
    ));
    await Promise.all(targets.map((stored) => this.cancelStored(stored, notifyWorker)));
  }

  async cancelGeneration(conversationId: string, generation: number): Promise<void> {
    const targets = [...this.pending.values()].filter((stored) => (
      stored.interaction.conversationId === conversationId && stored.generation === generation
    ));
    await Promise.all(targets.map((stored) => this.cancelStored(stored, false)));
  }

  private async cancelStored(stored: StoredInteraction, notifyWorker: boolean): Promise<void> {
    if (!this.pending.has(this.key(stored.interaction.conversationId, stored.interaction.id))) return;
    if (notifyWorker) {
      await this.transport.send(stored.interaction.conversationId, {
        type: 'extension_ui_response',
        id: stored.interaction.id,
        cancelled: true,
      }).catch(() => undefined);
    }
    this.finish(stored, 'cancelled');
  }

  private finish(
    stored: StoredInteraction,
    status: Exclude<ConversationInteraction['status'], 'pending'>,
  ): ConversationInteraction {
    this.pending.delete(this.key(stored.interaction.conversationId, stored.interaction.id));
    stored.untrack();
    const interaction = { ...stored.interaction, status };
    this.onChange(structuredClone(interaction));
    return interaction;
  }

  private key(conversationId: string, interactionId: string): string {
    return `${conversationId}\u0000${interactionId}`;
  }
}