makelore/tests/unit/pi-runtime-auth-recovery.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 { createCodingProjectAgent } from '../../electron/coding-projects/project-config';
import {
  createCodingProjectStore,
  createLocalCodingProject,
  createMemoryCodingProjectStorage,
} from '../../electron/coding-projects/project-store';
import { PiProcessError } from '../../electron/coding-runtime/pi/process-errors';
import { PiConversationRuntime } from '../../electron/coding-runtime/pi/runtime';
import { PiSessionRegistry } from '../../electron/coding-runtime/pi/session-registry';
import { PiWorkerPool, type PiConversationWorker } from '../../electron/coding-runtime/pi/worker-pool';
import type {
  PiRpcCommand,
  PiRpcEvent,
  PiRpcRequestOptions,
  PiRpcResponse,
} from '../../electron/coding-runtime/pi/rpc-client';

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

class AuthFailureWorker implements PiConversationWorker {
  readonly generation = 1;
  readonly requests: PiRpcCommand[] = [];
  constructor(readonly id: string) {}

  async request<T = unknown>(command: PiRpcCommand, _options?: PiRpcRequestOptions): Promise<PiRpcResponse<T>> {
    this.requests.push(structuredClone(command));
    if (command.type === 'prompt') {
      throw new PiProcessError('PI_RPC_RESPONSE_ERROR', '401 provider authentication failed');
    }
    const data = command.type === 'get_state'
      ? { sessionId: `session-${this.id}`, isStreaming: false, isCompacting: false }
      : command.type === 'get_entries'
        ? { entries: [], leafId: null }
        : command.type === 'get_session_stats'
          ? {
            contextUsage: { tokens: 0, contextWindow: 100_000, percent: 0 },
            tokens: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
          }
          : undefined;
    return {
      type: 'response',
      id: 'fake',
      success: true,
      ...(data === undefined ? {} : { data: data as T }),
    };
  }

  async send(_command: PiRpcCommand): Promise<void> {}

  subscribe(_listener: (event: PiRpcEvent) => void): () => void { return () => undefined; }
  subscribeInvalidation(_listener: (error: PiProcessError) => void): () => void { return () => undefined; }
  async stop() { return { mode: 'stdin-close' as const, code: 0, signal: null }; }
}

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

describe('Pi runtime Provider authentication recovery', () => {
  it('refreshes and reopens once, then exposes the second authentication failure without looping', async () => {
    const projectPath = await mkdtemp(path.join(tmpdir(), 'makelore-pi-auth-'));
    roots.push(projectPath);
    const projectStore = createCodingProjectStore(createMemoryCodingProjectStorage(), {
      createId: () => 'project-auth',
      now: () => NOW,
    });
    await createLocalCodingProject({ projectPath, now: NOW }, projectStore);
    await createCodingProjectAgent(projectPath, {
      id: 'agent-auth',
      avatarId: 'avatar-01',
      roleName: 'Implementer',
      name: 'Auth Agent',
      model: { accountId: 'account-auth', modelId: 'model-auth', thinkingLevel: 'medium' },
      modelResolution: 'resolved',
      responsibility: { mission: 'Implement', owns: [], boundaries: [], collaborators: [], principles: [] },
    }, { now: NOW });
    const store = createCodingConversationStore(projectPath, {
      createId: () => 'f47ac10b-58cc-4372-a567-0e02b2c3d479',
      now: () => NOW,
    });
    const conversation = await store.create({
      agentId: 'agent-auth',
      title: 'Auth Conversation',
      model: { accountId: 'account-auth', modelId: 'model-auth', thinkingLevel: 'medium' },
      modelResolution: 'resolved',
    });
    const workers: AuthFailureWorker[] = [];
    const pool = new PiWorkerPool({
      maxIdle: 2,
      openWorker: async ({ conversation: input, existingSession }) => {
        const worker = new AuthFailureWorker(`worker-${workers.length + 1}`);
        workers.push(worker);
        return {
          worker,
          session: existingSession ?? {
            piSessionId: `session-${input.conversationId}`,
            sessionKey: `key-${input.conversationId}`,
          },
        };
      },
    });
    const refreshCredential = vi.fn(async () => undefined);
    const runtime = new PiConversationRuntime({
      pool,
      registry: new PiSessionRegistry({ projectStore }),
      resolveModel: async () => { throw new Error('not used'); },
      refreshCredential,
      isAuthenticationError: (error) => (
        error instanceof PiProcessError && error.message.includes('401')
      ),
      createId: () => 'run-auth',
    });
    await runtime.prepare({
      conversationId: conversation.id,
      projectId: 'project-auth',
      agentId: 'agent-auth',
      title: conversation.title,
      model: { model: conversation.model, modelResolution: conversation.modelResolution },
    });

    await expect(runtime.prompt({
      clientRequestId: 'request-auth',
      conversationId: conversation.id,
      mode: 'prompt',
      text: 'Do not leak credentials',
      attachments: [],
    })).rejects.toMatchObject({
      code: 'PI_RPC_RESPONSE_ERROR',
    });

    await expect.poll(() => workers.length).toBe(2);
    await expect.poll(async () => (await runtime.getSnapshot(conversation.id)).run.status).toBe('error');
    expect(refreshCredential).toHaveBeenCalledTimes(1);
    expect(workers).toHaveLength(2);
    expect(workers.map((worker) => worker.requests.filter(({ type }) => type === 'prompt').length))
      .toEqual([1, 1]);
    expect(pool.getState(conversation.id)).toMatchObject({ state: 'idle', generation: 2 });
  });
});