makelore/electron/coding-runtime/pi/session-registry.ts

import {
  createCodingConversationStore,
  type CodingConversationV2,
  type PiSessionBinding,
} from '../../coding-projects/conversation-store';
import {
  readCodingProjectConfigV2,
  type CodingProjectAgentV2,
} from '../../coding-projects/project-config';
import type { CodingProjectStore } from '../../coding-projects/project-store';
import type {
  ConversationModelState,
  PrepareConversationInput,
} from '../contracts';
import { CodingRuntimeContractError } from '../runtime-errors';

export interface PiRegisteredConversation {
  projectPath: string;
  conversation: CodingConversationV2;
  agent: CodingProjectAgentV2;
  session: PiSessionBinding | null;
}

export interface PiSessionRegistryOptions {
  projectStore: CodingProjectStore;
  createConversationStore?: typeof createCodingConversationStore;
}

interface RegistryRecord extends PiRegisteredConversation {
  store: ReturnType<typeof createCodingConversationStore>;
}

function publicRecord(record: RegistryRecord): PiRegisteredConversation {
  return {
    projectPath: record.projectPath,
    conversation: structuredClone(record.conversation),
    agent: structuredClone(record.agent),
    session: record.session ? structuredClone(record.session) : null,
  };
}

function modelStateOf(conversation: CodingConversationV2): ConversationModelState {
  return {
    model: conversation.model ? structuredClone(conversation.model) : null,
    modelResolution: conversation.modelResolution,
  };
}

function sameModelState(left: ConversationModelState, right: ConversationModelState): boolean {
  return left.modelResolution === right.modelResolution
    && (left.model === null || right.model === null
      ? left.model === right.model
      : left.model.accountId === right.model.accountId
        && left.model.modelId === right.model.modelId
        && left.model.thinkingLevel === right.model.thinkingLevel
        && JSON.stringify(left.model.reasoningChoice) === JSON.stringify(right.model.reasoningChoice));
}

export class PiSessionRegistry {
  private readonly projectStore: CodingProjectStore;
  private readonly createConversationStore: typeof createCodingConversationStore;
  private readonly records = new Map<string, RegistryRecord>();
  private readonly prepareFlights = new Map<string, Promise<RegistryRecord>>();
  private readonly stores = new Map<string, ReturnType<typeof createCodingConversationStore>>();

  constructor(options: PiSessionRegistryOptions) {
    this.projectStore = options.projectStore;
    this.createConversationStore = options.createConversationStore ?? createCodingConversationStore;
  }

  async prepare(input: PrepareConversationInput): Promise<PiRegisteredConversation> {
    return publicRecord(await this.prepareRecord(input));
  }

  async ensureBinding(
    input: PrepareConversationInput,
    createBinding: () => Promise<PiSessionBinding>,
  ): Promise<PiRegisteredConversation> {
    const record = await this.prepareRecord(input);
    const conversation = await this.persistWrite(
      () => record.store.ensureSessionBinding(input.conversationId, createBinding),
    );
    record.conversation = conversation;
    record.session = {
      piSessionId: conversation.piSessionId as string,
      sessionKey: conversation.sessionKey as string,
    };
    return publicRecord(record);
  }

  async setModel(
    conversationId: string,
    model: ConversationModelState,
  ): Promise<ConversationModelState> {
    const record = this.records.get(conversationId);
    if (!record) throw new Error('Conversation is not registered');
    record.conversation = await this.persistWrite(
      () => record.store.setModelState(conversationId, model),
    );
    return modelStateOf(record.conversation);
  }

  forget(conversationId: string): void {
    this.records.delete(conversationId);
  }

  private prepareRecord(input: PrepareConversationInput): Promise<RegistryRecord> {
    const existing = this.records.get(input.conversationId);
    if (existing) return Promise.resolve(existing);
    const pending = this.prepareFlights.get(input.conversationId);
    if (pending) return pending;
    const flight = this.load(input).finally(() => {
      if (this.prepareFlights.get(input.conversationId) === flight) {
        this.prepareFlights.delete(input.conversationId);
      }
    });
    this.prepareFlights.set(input.conversationId, flight);
    return flight;
  }

  private async load(input: PrepareConversationInput): Promise<RegistryRecord> {
    const project = (await this.projectStore.listProjects())
      .find((candidate) => candidate.id === input.projectId);
    if (!project) throw new Error('Coding project does not exist');
    const configRead = await readCodingProjectConfigV2(project.path);
    if (configRead.status !== 'valid') throw new Error('Coding project configuration is unavailable');
    const agent = configRead.config.agents.find((candidate) => (
      candidate.id === input.agentId && candidate.enabled && !candidate.archivedAt
    ));
    if (!agent) throw new Error('Coding Agent does not exist');
    let store = this.stores.get(project.path);
    if (!store) {
      store = this.createConversationStore(project.path);
      this.stores.set(project.path, store);
    }
    const conversation = await store.get(input.conversationId);
    if (!conversation || conversation.agentId !== input.agentId) {
      throw new Error('Coding Conversation does not exist for the selected Agent');
    }
    if (!sameModelState(modelStateOf(conversation), input.model)) {
      throw new Error('Coding Conversation model metadata is stale');
    }
    const record: RegistryRecord = {
      projectPath: project.path,
      conversation,
      agent,
      session: conversation.piSessionId && conversation.sessionKey
        ? { piSessionId: conversation.piSessionId, sessionKey: conversation.sessionKey }
        : null,
      store,
    };
    this.records.set(input.conversationId, record);
    return record;
  }

  private async persistWrite<T>(operation: () => Promise<T>): Promise<T> {
    try {
      return await operation();
    } catch (error) {
      if (error instanceof CodingRuntimeContractError) throw error;
      throw new CodingRuntimeContractError(
        'CODING_STORAGE_WRITE_FAILED',
        'Coding Conversation state could not be persisted',
        true,
      );
    }
  }
}