makelore/tests/unit/pi-session-registry.test.ts

// @vitest-environment node

import { mkdtemp, rm } from 'node:fs/promises';
import { tmpdir } from 'node:os';
import path from 'node:path';
import { afterEach, describe, expect, it, vi } from 'vitest';
import { createCodingConversationStore } from '../../electron/coding-projects/conversation-store';
import { CodingProjectService } from '../../electron/coding-projects/project-service';
import { createCodingProjectAgent } from '../../electron/coding-projects/project-config';
import {
  createCodingProjectStore,
  createLocalCodingProject,
  createMemoryCodingProjectStorage,
} from '../../electron/coding-projects/project-store';
import { PiSessionRegistry } from '../../electron/coding-runtime/pi/session-registry';

const roots: string[] = [];
const NOW = '2026-08-22T14:00:00.000Z';

afterEach(async () => {
  await Promise.all(roots.splice(0).map((root) => rm(root, { recursive: true, force: true })));
});

describe('Pi session registry', () => {
  it('persists one binding and target-only model state across registry reopen', async () => {
    const projectPath = await mkdtemp(path.join(tmpdir(), 'makelore-pi-registry-'));
    roots.push(projectPath);
    const projectStore = createCodingProjectStore(createMemoryCodingProjectStorage(), {
      createId: () => 'project-a',
      now: () => NOW,
    });
    const { project } = await createLocalCodingProject({ projectPath, now: NOW }, projectStore);
    await createCodingProjectAgent(projectPath, {
      id: 'agent-a',
      avatarId: 'avatar-01',
      roleName: 'Implementer',
      name: 'Agent A',
      model: { accountId: 'account-a', modelId: 'model-a', thinkingLevel: 'medium' },
      modelResolution: 'resolved',
      responsibility: { mission: 'Implement', owns: [], boundaries: [], collaborators: [], principles: [] },
      prompt: 'Managed prompt',
      skillIds: ['tdd'],
    }, { now: NOW });
    const conversations = createCodingConversationStore(projectPath, {
      createId: () => 'f47ac10b-58cc-4372-a567-0e02b2c3d479',
      now: () => NOW,
    });
    const created = await conversations.create({
      agentId: 'agent-a',
      title: 'Conversation A',
      model: { accountId: 'account-a', modelId: 'model-a', thinkingLevel: 'medium' },
      modelResolution: 'resolved',
    });
    const input = {
      conversationId: created.id,
      projectId: 'project-a',
      agentId: 'agent-a',
      title: created.title,
      model: { model: created.model, modelResolution: created.modelResolution },
    } as const;
    const registry = new PiSessionRegistry({ projectStore });
    const createBinding = vi.fn(async () => ({
      piSessionId: 'pi-session-a',
      sessionKey: 'session-key-a',
    }));

    const [first, duplicate] = await Promise.all([
      registry.ensureBinding(input, createBinding),
      registry.ensureBinding(input, createBinding),
    ]);
    expect(createBinding).toHaveBeenCalledTimes(1);
    expect(first.session).toEqual({ piSessionId: 'pi-session-a', sessionKey: 'session-key-a' });
    expect(duplicate).toEqual(first);

    await registry.setModel(created.id, {
      model: { accountId: 'account-b', modelId: 'model-b', thinkingLevel: 'high' },
      modelResolution: 'resolved',
    });
    const reopened = await new PiSessionRegistry({ projectStore }).prepare({
      ...input,
      model: {
        model: { thinkingLevel: 'high', modelId: 'model-b', accountId: 'account-b' },
        modelResolution: 'resolved',
      },
    });
    expect(reopened).toMatchObject({
      projectPath: project.path,
      agent: { id: 'agent-a', prompt: 'Managed prompt', skillIds: ['tdd'] },
      conversation: {
        id: created.id,
        model: { accountId: 'account-b', modelId: 'model-b', thinkingLevel: 'high' },
      },
      session: { piSessionId: 'pi-session-a', sessionKey: 'session-key-a' },
    });
  });

  it('serializes first bindings for different Conversations in one project', async () => {
    const projectPath = await mkdtemp(path.join(tmpdir(), 'makelore-pi-registry-concurrent-'));
    roots.push(projectPath);
    const projectStore = createCodingProjectStore(createMemoryCodingProjectStorage(), {
      createId: () => 'project-concurrent',
      now: () => NOW,
    });
    await createLocalCodingProject({ projectPath, now: NOW }, projectStore);
    await createCodingProjectAgent(projectPath, {
      id: 'agent-a',
      avatarId: 'avatar-01',
      roleName: 'Implementer',
      name: 'Agent A',
      model: { accountId: 'account-a', modelId: 'model-a', thinkingLevel: 'medium' },
      modelResolution: 'resolved',
      responsibility: { mission: 'Implement', owns: [], boundaries: [], collaborators: [], principles: [] },
    }, { now: NOW });
    const conversationIds = [
      'f47ac10b-58cc-4372-a567-0e02b2c3d482',
      'f47ac10b-58cc-4372-a567-0e02b2c3d483',
    ];
    const conversations = createCodingConversationStore(projectPath, {
      createId: () => conversationIds.shift() as string,
      now: () => NOW,
    });
    const left = await conversations.create({
      agentId: 'agent-a',
      title: 'Left',
      model: { accountId: 'account-a', modelId: 'model-a', thinkingLevel: 'medium' },
      modelResolution: 'resolved',
    });
    const right = await conversations.create({
      agentId: 'agent-a',
      title: 'Right',
      model: { accountId: 'account-a', modelId: 'model-a', thinkingLevel: 'medium' },
      modelResolution: 'resolved',
    });
    let storeCreations = 0;
    const registry = new PiSessionRegistry({
      projectStore,
      createConversationStore: (candidatePath, options) => {
        storeCreations += 1;
        return createCodingConversationStore(candidatePath, options);
      },
    });
    const prepareInput = (conversation: typeof left) => ({
      conversationId: conversation.id,
      projectId: 'project-concurrent',
      agentId: 'agent-a',
      title: conversation.title,
      model: { model: conversation.model, modelResolution: conversation.modelResolution },
    });

    await Promise.all([left, right].map((conversation) => registry.ensureBinding(
      prepareInput(conversation),
      async () => ({
        piSessionId: `session-${conversation.id}`,
        sessionKey: `key-${conversation.id}`,
      }),
    )));

    const reopened = await conversations.read();
    expect(storeCreations).toBe(1);
    expect(reopened.conversations).toEqual(expect.arrayContaining([
      expect.objectContaining({
        id: left.id,
        piSessionId: `session-${left.id}`,
        sessionKey: `key-${left.id}`,
      }),
      expect.objectContaining({
        id: right.id,
        piSessionId: `session-${right.id}`,
        sessionKey: `key-${right.id}`,
      }),
    ]));
  });

  it('shares one project mutation queue with ProjectService metadata writes', async () => {
    const projectPath = await mkdtemp(path.join(tmpdir(), 'makelore-pi-shared-conversations-'));
    roots.push(projectPath);
    const projectStore = createCodingProjectStore(createMemoryCodingProjectStorage(), {
      createId: () => 'project-shared',
      now: () => NOW,
    });
    await createLocalCodingProject({ projectPath, now: NOW }, projectStore);
    await createCodingProjectAgent(projectPath, {
      id: 'agent-a',
      avatarId: 'avatar-01',
      roleName: 'Implementer',
      name: 'Agent A',
      model: { accountId: 'account-a', modelId: 'model-a', thinkingLevel: 'medium' },
      modelResolution: 'resolved',
      responsibility: { mission: 'Implement', owns: [], boundaries: [], collaborators: [], principles: [] },
    }, { now: NOW });
    const conversationIds = [
      'f47ac10b-58cc-4372-a567-0e02b2c3d484',
      'f47ac10b-58cc-4372-a567-0e02b2c3d485',
    ];
    const sharedStore = createCodingConversationStore(projectPath, {
      createId: () => conversationIds.shift() as string,
      now: () => NOW,
    });
    const left = await sharedStore.create({
      agentId: 'agent-a',
      title: 'Left',
      model: { accountId: 'account-a', modelId: 'model-a', thinkingLevel: 'medium' },
      modelResolution: 'resolved',
    });
    const sharedProvider = vi.fn(() => sharedStore);
    const projects = new CodingProjectService(projectStore, {
      createConversationStore: sharedProvider,
    });
    const registry = new PiSessionRegistry({
      projectStore,
      createConversationStore: sharedProvider,
    });

    await Promise.all([
      registry.ensureBinding({
        conversationId: left.id,
        projectId: 'project-shared',
        agentId: 'agent-a',
        title: left.title,
        model: { model: left.model, modelResolution: left.modelResolution },
      }, async () => ({ piSessionId: 'pi-left', sessionKey: 'key-left' })),
      projects.conversationStore(projectPath).create({
        agentId: 'agent-a',
        title: 'Right',
        model: { accountId: 'account-a', modelId: 'model-a', thinkingLevel: 'medium' },
        modelResolution: 'resolved',
      }),
    ]);

    const persisted = await sharedStore.read();
    expect(projects.conversationStore(projectPath)).toBe(sharedStore);
    expect(persisted.conversations).toEqual(expect.arrayContaining([
      expect.objectContaining({
        id: left.id,
        piSessionId: 'pi-left',
        sessionKey: 'key-left',
      }),
      expect.objectContaining({ title: 'Right', agentId: 'agent-a' }),
    ]));
  });

  it('maps first session binding persistence failure to the stable storage error', async () => {
    const projectPath = await mkdtemp(path.join(tmpdir(), 'makelore-pi-registry-binding-failure-'));
    roots.push(projectPath);
    const projectStore = createCodingProjectStore(createMemoryCodingProjectStorage(), {
      createId: () => 'project-binding-failure',
      now: () => NOW,
    });
    await createLocalCodingProject({ projectPath, now: NOW }, projectStore);
    await createCodingProjectAgent(projectPath, {
      id: 'agent-a',
      avatarId: 'avatar-01',
      roleName: 'Implementer',
      name: 'Agent A',
      model: { accountId: 'account-a', modelId: 'model-a', thinkingLevel: 'medium' },
      modelResolution: 'resolved',
      responsibility: { mission: 'Implement', owns: [], boundaries: [], collaborators: [], principles: [] },
    }, { now: NOW });
    const conversations = createCodingConversationStore(projectPath, {
      createId: () => 'f47ac10b-58cc-4372-a567-0e02b2c3d480',
      now: () => NOW,
    });
    const created = await conversations.create({
      agentId: 'agent-a',
      title: 'Conversation A',
      model: { accountId: 'account-a', modelId: 'model-a', thinkingLevel: 'medium' },
      modelResolution: 'resolved',
    });
    const registry = new PiSessionRegistry({
      projectStore,
      createConversationStore: () => conversations,
    });
    const input = {
      conversationId: created.id,
      projectId: 'project-binding-failure',
      agentId: 'agent-a',
      title: created.title,
      model: { model: created.model, modelResolution: created.modelResolution },
    } as const;
    await registry.prepare(input);
    vi.spyOn(conversations, 'ensureSessionBinding').mockRejectedValueOnce(new Error('disk full'));

    await expect(registry.ensureBinding(input, async () => ({
      piSessionId: 'pi-session-a',
      sessionKey: 'session-key-a',
    }))).rejects.toMatchObject({
      publicError: { code: 'CODING_STORAGE_WRITE_FAILED', recoverable: true },
    });
  });

  it('maps model metadata persistence failure to the stable storage error', async () => {
    const projectPath = await mkdtemp(path.join(tmpdir(), 'makelore-pi-registry-model-failure-'));
    roots.push(projectPath);
    const projectStore = createCodingProjectStore(createMemoryCodingProjectStorage(), {
      createId: () => 'project-model-failure',
      now: () => NOW,
    });
    await createLocalCodingProject({ projectPath, now: NOW }, projectStore);
    await createCodingProjectAgent(projectPath, {
      id: 'agent-a',
      avatarId: 'avatar-01',
      roleName: 'Implementer',
      name: 'Agent A',
      model: { accountId: 'account-a', modelId: 'model-a', thinkingLevel: 'medium' },
      modelResolution: 'resolved',
      responsibility: { mission: 'Implement', owns: [], boundaries: [], collaborators: [], principles: [] },
    }, { now: NOW });
    const conversations = createCodingConversationStore(projectPath, {
      createId: () => 'f47ac10b-58cc-4372-a567-0e02b2c3d481',
      now: () => NOW,
    });
    const created = await conversations.create({
      agentId: 'agent-a',
      title: 'Conversation A',
      model: { accountId: 'account-a', modelId: 'model-a', thinkingLevel: 'medium' },
      modelResolution: 'resolved',
    });
    const registry = new PiSessionRegistry({
      projectStore,
      createConversationStore: () => conversations,
    });
    await registry.prepare({
      conversationId: created.id,
      projectId: 'project-model-failure',
      agentId: 'agent-a',
      title: created.title,
      model: { model: created.model, modelResolution: created.modelResolution },
    });
    vi.spyOn(conversations, 'setModelState').mockRejectedValueOnce(new Error('disk full'));

    await expect(registry.setModel(created.id, {
      model: { accountId: 'account-b', modelId: 'model-b', thinkingLevel: 'high' },
      modelResolution: 'resolved',
    })).rejects.toMatchObject({
      publicError: { code: 'CODING_STORAGE_WRITE_FAILED', recoverable: true },
    });
  });
});