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 {
  CodingProviderCredentialRefreshError,
  isCodingProviderAuthenticationError,
} from '../../electron/api/coding-provider-auth';
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 { PiProviderConfigError } from '../../electron/coding-runtime/pi/provider-config';
import { PiConversationRuntime } from '../../electron/coding-runtime/pi/runtime';
import { PiSessionRegistry } from '../../electron/coding-runtime/pi/session-registry';
import {
  PiWorkerPool,
  type PiConversationWorker,
  type PiWorkerPoolOptions,
} 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 })));
});

async function setupAuthRuntime(
  refreshCredential: (accountId: string) => Promise<void>,
  options: {
    openWorker?: PiWorkerPoolOptions['openWorker'];
    prepare?: boolean;
  } = {},
) {
  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 defaultOpenWorker: PiWorkerPoolOptions['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 pool = new PiWorkerPool({
    maxIdle: 2,
    openWorker: options.openWorker ?? defaultOpenWorker,
  });
  const runtime = new PiConversationRuntime({
    pool,
    registry: new PiSessionRegistry({ projectStore }),
    resolveModel: async () => { throw new Error('not used'); },
    refreshCredential,
    isAuthenticationError: isCodingProviderAuthenticationError,
    createId: () => 'run-auth',
  });
  const input = {
    conversationId: conversation.id,
    projectId: 'project-auth',
    agentId: 'agent-auth',
    title: conversation.title,
    model: { model: conversation.model, modelResolution: conversation.modelResolution },
  };
  if (options.prepare !== false) await runtime.prepare(input);
  return { conversation, input, pool, runtime, workers };
}

describe('Pi runtime Provider authentication recovery', () => {
  it('maps a missing proxy token to Provider auth-required after one refresh retry', async () => {
    const refreshCredential = vi.fn(async () => undefined);
    const openWorker = vi.fn<PiWorkerPoolOptions['openWorker']>(async () => {
      throw new PiProviderConfigError('PROVIDER_AUTH_REQUIRED', 'Provider credential is unavailable');
    });
    const { input, runtime } = await setupAuthRuntime(refreshCredential, {
      openWorker,
      prepare: false,
    });

    await expect(runtime.prepare(input)).rejects.toMatchObject({
      publicError: {
        code: 'CODING_PROVIDER_AUTH_REQUIRED',
        recoverable: true,
      },
    });
    expect(refreshCredential).toHaveBeenCalledTimes(1);
    expect(openWorker).toHaveBeenCalledTimes(2);
  });

  it('refreshes and reopens once, then exposes the second authentication failure without looping', async () => {
    const refreshCredential = vi.fn(async () => undefined);
    const { conversation, pool, runtime, workers } = await setupAuthRuntime(refreshCredential);

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

    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 });
  });

  it('projects a rejected credential refresh as authentication required', async () => {
    const refreshCredential = vi.fn(async () => {
      throw new CodingProviderCredentialRefreshError('Provider credential refresh failed');
    });
    const { conversation, runtime } = await setupAuthRuntime(refreshCredential);

    await expect(runtime.prompt({
      clientRequestId: 'request-auth-refresh-rejected',
      conversationId: conversation.id,
      mode: 'prompt',
      text: 'Do not leak refresh failures',
      attachments: [],
    })).rejects.toMatchObject({
      publicError: {
        code: 'CODING_PROVIDER_AUTH_REQUIRED',
        recoverable: true,
      },
    });
    expect(refreshCredential).toHaveBeenCalledTimes(1);
  });
});