makelore/electron/coding-runtime/in-memory-conversation-runtime.ts

import type {
  CodingConversationRuntime,
  CodingRuntimeCommand,
  CodingRuntimeDiagnostics,
  CodingRuntimeDisposeReason,
  ConversationModelState,
  ConversationInteraction,
  ConversationInteractionResponse,
  ConversationPatch,
  ConversationPatchEnvelope,
  ConversationRuntimeState,
  ConversationSnapshot,
  ForkConversationInput,
  ForkResult,
  PrepareConversationInput,
  ProductModelRef,
  PromptAcceptance,
  PromptConversationInput,
  QueueAcceptance,
  QueueMessageInput,
  SetConversationModelInput,
  SetThinkingLevelInput,
} from './contracts';
import { CodingRuntimeContractError } from './runtime-errors';
export { CodingRuntimeContractError } from './runtime-errors';
import {
  createConversationReducerState,
  reduceConversationPatch,
  type ConversationReducerState,
} from './conversation-reducer';

export interface InMemoryConversationRuntimeOptions {
  snapshots?: ConversationSnapshot[];
  commands?: CodingRuntimeCommand[];
  now?: () => number;
  createId?: (kind: 'run' | 'node' | 'queue' | 'compaction') => string;
}

function clone<T>(value: T): T {
  return structuredClone(value);
}

function emptySnapshot(input: PrepareConversationInput): ConversationSnapshot {
  return {
    schemaVersion: 1,
    conversation: {
      id: input.conversationId,
      projectId: input.projectId,
      agentId: input.agentId,
      title: input.title,
      model: clone(input.model),
    },
    nodes: [],
    run: { status: 'idle' },
    queue: { items: [] },
    context: {
      usedTokens: 0,
      contextWindow: 0,
      compaction: 'idle',
    },
    pendingInteractions: [],
    worker: {
      status: 'ready',
      generation: 0,
    },
    cursor: {
      workerGeneration: 0,
      seq: 0,
    },
  };
}

export class InMemoryConversationRuntime implements CodingConversationRuntime {
  private readonly states = new Map<string, ConversationReducerState>();
  private readonly listeners = new Set<(patch: ConversationPatchEnvelope) => void>();
  private readonly promptAcceptances = new Map<string, PromptAcceptance>();
  private readonly queueAcceptances = new Map<string, QueueAcceptance>();
  private readonly now: () => number;
  private readonly commands: CodingRuntimeCommand[];
  private readonly createId: InMemoryConversationRuntimeOptions['createId'];
  private nextId = 0;
  private providerRevision = 1;
  private resourcesRevision = 1;

  constructor(options: InMemoryConversationRuntimeOptions = {}) {
    this.now = options.now ?? Date.now;
    this.commands = clone(options.commands ?? []);
    this.createId = options.createId;
    for (const snapshot of options.snapshots ?? []) {
      const state = createConversationReducerState(snapshot);
      if (!state.snapshot || state.invalidation) {
        throw new CodingRuntimeContractError(
          'CODING_SESSION_UNREADABLE',
          'Initial Conversation snapshot is invalid',
          true,
        );
      }
      this.states.set(snapshot.conversation.id, state);
    }
  }

  private id(kind: 'run' | 'node' | 'queue' | 'compaction'): string {
    if (this.createId) return this.createId(kind);
    this.nextId += 1;
    return `${kind}-${this.nextId}`;
  }

  private state(conversationId: string): ConversationReducerState {
    const state = this.states.get(conversationId);
    if (!state?.snapshot) {
      throw new CodingRuntimeContractError(
        'CODING_CONVERSATION_NOT_FOUND',
        'Conversation is not prepared',
        true,
      );
    }
    return state;
  }

  private snapshot(conversationId: string): ConversationSnapshot {
    return this.state(conversationId).snapshot as ConversationSnapshot;
  }

  private replaceSnapshot(snapshot: ConversationSnapshot): void {
    const state = createConversationReducerState(snapshot);
    if (!state.snapshot || state.invalidation) {
      throw new CodingRuntimeContractError(
        'CODING_SESSION_UNREADABLE',
        'Conversation snapshot replacement is invalid',
        true,
      );
    }
    this.states.set(snapshot.conversation.id, state);
  }

  private emit(conversationId: string, patch: ConversationPatch, runId?: string): void {
    const state = this.state(conversationId);
    const snapshot = state.snapshot as ConversationSnapshot;
    const envelope: ConversationPatchEnvelope = {
      conversationId,
      workerGeneration: snapshot.cursor.workerGeneration,
      ...(runId ? { runId } : {}),
      seq: snapshot.cursor.seq + 1,
      at: this.now(),
      patch: clone(patch),
    };
    const next = reduceConversationPatch(state, envelope);
    if (next.invalidation) {
      throw new CodingRuntimeContractError(
        'CODING_RUNTIME_PROTOCOL_ERROR',
        next.invalidation.reason,
        true,
      );
    }
    this.states.set(conversationId, next);
    for (const listener of this.listeners) listener(clone(envelope));
  }

  private runtimeState(conversationId: string): ConversationRuntimeState {
    const snapshot = this.snapshot(conversationId);
    return {
      conversationId,
      status: snapshot.worker.status,
      workerGeneration: snapshot.worker.generation,
      ...(snapshot.worker.error ? { error: clone(snapshot.worker.error) } : {}),
    };
  }

  async prepare(input: PrepareConversationInput): Promise<ConversationRuntimeState> {
    const existing = this.states.get(input.conversationId)?.snapshot;
    if (!existing) this.replaceSnapshot(emptySnapshot(input));
    return this.runtimeState(input.conversationId);
  }

  async getSnapshot(conversationId: string): Promise<ConversationSnapshot> {
    return clone(this.snapshot(conversationId));
  }

  async prompt(input: PromptConversationInput): Promise<PromptAcceptance> {
    if (input.mode === 'steer' || input.mode === 'follow-up') {
      const acceptance = input.mode === 'steer'
        ? await this.steer(input)
        : await this.followUp(input);
      return {
        accepted: true,
        conversationId: input.conversationId,
        clientRequestId: input.clientRequestId,
        runId: this.snapshot(input.conversationId).run.runId ?? this.id('run'),
        mode: input.mode,
        queuePosition: acceptance.queuePosition,
      };
    }
    const key = `${input.conversationId}:${input.clientRequestId}`;
    const prior = this.promptAcceptances.get(key);
    if (prior) return clone(prior);

    const snapshot = this.snapshot(input.conversationId);
    if (snapshot.worker.status !== 'ready') {
      throw new CodingRuntimeContractError(
        'CODING_RUNTIME_START_FAILED',
        'Conversation worker is not ready',
        true,
      );
    }

    const runId = this.id('run');
    const messageId = this.id('node');
    this.emit(input.conversationId, {
      op: 'message.upsert',
      node: {
        kind: 'message',
        id: messageId,
        clientRequestId: input.clientRequestId,
        role: 'user',
        status: 'optimistic',
        blocks: [
          ...(input.text ? [{
            kind: 'text' as const,
            id: `${messageId}:text`,
            text: input.text,
            status: 'complete' as const,
          }] : []),
          ...input.attachments.map(({ attachmentId }, index) => ({
            kind: 'image' as const,
            id: `${messageId}:attachment:${index}`,
            attachmentId,
            mime: 'application/octet-stream',
          })),
        ],
      },
    }, runId);
    this.emit(input.conversationId, {
      op: 'run.state',
      run: {
        status: 'running',
        runId,
        mode: input.mode,
        startedAt: this.now(),
      },
    }, runId);

    const acceptance: PromptAcceptance = {
      accepted: true,
      conversationId: input.conversationId,
      clientRequestId: input.clientRequestId,
      runId,
      mode: input.mode,
    };
    this.promptAcceptances.set(key, acceptance);
    return clone(acceptance);
  }

  private async queue(
    mode: 'steer' | 'follow-up',
    input: QueueMessageInput,
  ): Promise<QueueAcceptance> {
    const key = `${input.conversationId}:${mode}:${input.clientRequestId}`;
    const prior = this.queueAcceptances.get(key);
    if (prior) return clone(prior);

    const snapshot = this.snapshot(input.conversationId);
    const queuePosition = snapshot.queue.items.length + 1;
    this.emit(input.conversationId, {
      op: 'queue.replace',
      queue: {
        items: [
          ...snapshot.queue.items,
          {
            id: this.id('queue'),
            clientRequestId: input.clientRequestId,
            mode,
            text: input.text,
            attachmentIds: input.attachments.map(({ attachmentId }) => attachmentId),
          },
        ],
      },
    });

    const acceptance: QueueAcceptance = {
      accepted: true,
      conversationId: input.conversationId,
      clientRequestId: input.clientRequestId,
      mode,
      queuePosition,
    };
    this.queueAcceptances.set(key, acceptance);
    return clone(acceptance);
  }

  async steer(input: QueueMessageInput): Promise<QueueAcceptance> {
    return this.queue('steer', input);
  }

  async followUp(input: QueueMessageInput): Promise<QueueAcceptance> {
    return this.queue('follow-up', input);
  }

  async abort(conversationId: string): Promise<void> {
    const current = this.snapshot(conversationId).run;
    this.emit(conversationId, {
      op: 'run.state',
      run: {
        status: 'idle',
        ...(current.runId ? { runId: current.runId } : {}),
        ...(current.mode ? { mode: current.mode } : {}),
        ...(current.startedAt !== undefined ? { startedAt: current.startedAt } : {}),
        settledAt: this.now(),
        terminalReason: 'aborted',
      },
    }, current.runId);
  }

  async validateModel(model: ProductModelRef): Promise<ProductModelRef> {
    if (!model.accountId.trim() || !model.modelId.trim()) {
      throw new CodingRuntimeContractError(
        'CODING_MODEL_UNAVAILABLE',
        'The selected model is unavailable',
        true,
      );
    }
    return clone(model);
  }

  async setModel(input: SetConversationModelInput): Promise<ConversationModelState> {
    const snapshot = this.snapshot(input.conversationId);
    const thinkingLevel = snapshot.conversation.model.model?.thinkingLevel ?? 'off';
    const model: ConversationModelState = {
      model: {
        accountId: input.accountId,
        modelId: input.modelId,
        thinkingLevel,
      },
      modelResolution: 'resolved',
    };
    this.replaceSnapshot({
      ...snapshot,
      conversation: { ...snapshot.conversation, model },
    });
    return clone(model);
  }

  async setThinking(input: SetThinkingLevelInput): Promise<ConversationModelState> {
    const snapshot = this.snapshot(input.conversationId);
    if (!snapshot.conversation.model.model) {
      throw new CodingRuntimeContractError(
        'CODING_MIGRATION_MODEL_REQUIRED',
        'Conversation model must be selected first',
        true,
      );
    }
    const model: ConversationModelState = {
      model: {
        ...snapshot.conversation.model.model,
        thinkingLevel: input.thinkingLevel,
      },
      modelResolution: 'resolved',
    };
    this.replaceSnapshot({
      ...snapshot,
      conversation: { ...snapshot.conversation, model },
    });
    return clone(model);
  }

  async compact(conversationId: string): Promise<void> {
    const before = this.snapshot(conversationId);
    const compactionId = this.id('compaction');
    const runId = before.run.runId ?? this.id('run');
    const willRetry = before.run.status !== 'idle';
    this.emit(conversationId, {
      op: 'compaction.upsert',
      node: {
        kind: 'compaction',
        id: compactionId,
        runId,
        source: 'manual',
        status: 'running',
        willRetry,
      },
    }, runId);
    this.emit(conversationId, {
      op: 'context.replace',
      context: {
        ...before.context,
        compaction: 'running',
        lastCompactionId: compactionId,
      },
    }, runId);
    this.emit(conversationId, {
      op: 'compaction.upsert',
      node: {
        kind: 'compaction',
        id: compactionId,
        runId,
        source: 'manual',
        status: 'complete',
        willRetry,
      },
    }, runId);
    this.emit(conversationId, {
      op: 'context.replace',
      context: {
        ...before.context,
        compaction: 'idle',
        lastCompactionId: compactionId,
      },
    }, runId);
  }

  async fork(input: ForkConversationInput): Promise<ForkResult> {
    const source = this.snapshot(input.sourceConversationId);
    let nodes = source.nodes;
    if (input.sourceEntryId) {
      const index = source.nodes.findIndex(
        (node) => node.kind === 'message' && node.sourceEntryId === input.sourceEntryId,
      );
      if (index < 0) {
        throw new CodingRuntimeContractError(
          'CODING_SESSION_UNREADABLE',
          'Fork source entry is not on the active product path',
          true,
        );
      }
      nodes = source.nodes.slice(0, index + 1);
    }
    const snapshot: ConversationSnapshot = {
      ...emptySnapshot(input.conversation),
      nodes: clone(nodes),
      context: clone(source.context),
      cursor: {
        workerGeneration: 0,
        seq: 0,
        ...(input.sourceEntryId ? { leafEntryId: input.sourceEntryId } : {}),
      },
      worker: {
        status: 'stopped',
        generation: 0,
      },
    };
    this.replaceSnapshot(snapshot);
    return { conversationId: input.conversation.conversationId, snapshot: clone(snapshot) };
  }

  async recover(conversationId: string): Promise<ConversationRuntimeState> {
    const snapshot = this.snapshot(conversationId);
    const generation = snapshot.cursor.workerGeneration + 1;
    this.replaceSnapshot({
      ...snapshot,
      worker: { status: 'ready', generation },
      run: { status: 'idle' },
      pendingInteractions: [],
      cursor: {
        ...snapshot.cursor,
        workerGeneration: generation,
        seq: 0,
      },
    });
    return this.runtimeState(conversationId);
  }

  async dispose(conversationId: string, _reason: CodingRuntimeDisposeReason): Promise<void> {
    if (!this.states.has(conversationId)) return;
    const snapshot = this.snapshot(conversationId);
    this.replaceSnapshot({
      ...snapshot,
      worker: {
        status: 'stopped',
        generation: snapshot.cursor.workerGeneration,
      },
      run: { status: 'idle' },
      pendingInteractions: [],
    });
  }

  async listCommands(conversationId: string): Promise<CodingRuntimeCommand[]> {
    return this.states.has(conversationId) ? clone(this.commands) : [];
  }

  async listInteractions(conversationId?: string): Promise<ConversationInteraction[]> {
    return [...this.states.values()].flatMap((state) => {
      const snapshot = state.snapshot;
      if (!snapshot || (conversationId && snapshot.conversation.id !== conversationId)) return [];
      return clone(snapshot.pendingInteractions);
    });
  }

  async respondInteraction(
    conversationId: string,
    response: ConversationInteractionResponse,
  ): Promise<void> {
    const interaction = this.snapshot(conversationId).pendingInteractions.find(
      ({ id }) => id === response.interactionId,
    );
    if (!interaction) {
      throw new CodingRuntimeContractError(
        'CODING_RUNTIME_PROTOCOL_ERROR',
        'Conversation interaction is not pending',
        true,
      );
    }
    this.emit(conversationId, { op: 'interaction.remove', interactionId: interaction.id }, interaction.runId);
  }

  getDiagnostics(): CodingRuntimeDiagnostics {
    return {
      revision: {
        provider: this.providerRevision,
        resources: this.resourcesRevision,
      },
      workers: [...this.states.values()].flatMap((state) => {
        const snapshot = state.snapshot;
        if (!snapshot) return [];
        const runtimeState = snapshot.worker.status === 'error'
          ? 'crashed' as const
          : snapshot.worker.status === 'starting' || snapshot.worker.status === 'recovering'
            ? 'spawning' as const
            : snapshot.worker.status === 'stopped'
              ? 'idle' as const
              : snapshot.run.status === 'queued'
                ? 'queued' as const
                : snapshot.run.status === 'running'
                  ? 'running' as const
                  : 'ready' as const;
        return [{
          conversationId: snapshot.conversation.id,
          generation: snapshot.worker.generation,
          state: runtimeState,
          stage: runtimeState === 'crashed'
            ? 'failed' as const
            : runtimeState === 'spawning'
              ? 'starting' as const
              : runtimeState === 'queued'
                ? 'queued' as const
                : runtimeState === 'running'
                  ? 'running' as const
                  : 'idle' as const,
        }];
      }),
    };
  }

  markProviderStale(): void {
    this.providerRevision += 1;
  }

  markResourcesStale(): void {
    this.resourcesRevision += 1;
  }

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