open-design/e2e/lib/amr.ts

import { chmod, mkdir, rm, writeFile } from 'node:fs/promises';
import { dirname, join } from 'node:path';

export type FakeVelaOptions = {
  assistantText?: string;
  endpoints?: VelaEndpoints;
  failAuthAtPrompt?: boolean;
  failAuthAtPromptOnce?: boolean;
  failBalanceAtPrompt?: boolean;
  failBalanceAtPromptOnce?: boolean;
  failModelListInvalidApiKey?: boolean;
  fundingCheck?: { status: 'available' | 'insufficient' | 'unknown'; workspaceMemberId: string };
  requireLoginConfig?: boolean;
  requireSetModel?: boolean;
  sessionId?: string;
};

export type VelaEndpoints = {
  apiUrl: string;
  linkUrl: string;
};

const DEFAULT_ASSISTANT_TEXT = 'Hello from the e2e fake vela.';
/**
 * What the real `vela` CLI prints when the API rejects its runtime key mid-run.
 *
 * vela's API answers an unauthenticated request with `{"error":"unauthenticated"}`
 * under HTTP 401 (`services/api/src/app.ts:723`, `:792`, `:863`;
 * `runtime-keys.ts:51`), and the CLI renders that through `client.ParseAPIError`
 * as `API request failed with status %d: %s`
 * (`apps/cli/internal/client/client.go:42`).
 *
 * It has to be a shape vela actually sends. `classifyAmrAccountFailure`
 * (`apps/daemon/src/integrations/vela-errors.ts`) only claims the AMR-specific
 * `AMR_AUTH_REQUIRED` on evidence that identifies the AMR credential — a whole
 * `unauthenticated` / `auth_required` code, `not logged in` alongside vela/amr,
 * or the sign-in `session` noun — because every run's failure text also carries
 * `gh`, `npm`, `curl` and MCP output, whose sign-in prose belongs to somebody
 * else. Invented sign-in wording is deliberately not enough, so a fabricated
 * message here would assert a path production never takes.
 */
const PROMPT_AUTH_FAILURE_MESSAGE = 'API request failed with status 401: unauthenticated';
const PRESET_MODELS_JSON = JSON.stringify({
  source: 'preset',
  data: [
    { id: 'glm-5' },
  ],
});
const REMOTE_MODELS_JSON = JSON.stringify({
  source: 'remote',
  data: [
    { id: 'glm-5' },
  ],
});

export async function writeFakeVelaBin(root: string, options: FakeVelaOptions = {}): Promise<string> {
  await mkdir(root, { recursive: true });
  const bin = join(root, 'vela');
  await writeFile(bin, renderFakeVelaScript(options), 'utf8');
  await chmod(bin, 0o755);
  return bin;
}

export async function seedVelaLoginConfig(
  homeDir: string,
  options: {
    endpoints?: VelaEndpoints;
    profile?: string;
    email?: string;
    runtimeKey?: string;
    controlKey?: string;
  } = {},
): Promise<string> {
  const profile = options.profile ?? 'local';
  const endpoints = options.endpoints ?? defaultVelaEndpoints();
  const configDir = join(homeDir, '.amr');
  const file = join(configDir, 'config.json');
  await mkdir(configDir, { recursive: true });
  await writeFile(
    file,
    JSON.stringify(
      {
        profiles: {
          [profile]: {
            runtimeKey: options.runtimeKey ?? 'fake-runtime-key',
            controlKey: options.controlKey ?? 'fake-control-key',
            apiUrl: endpoints.apiUrl,
            linkUrl: endpoints.linkUrl,
            user: {
              id: 'fake-user-id',
              email: options.email ?? 'e2e@example.com',
              plan: 'free',
            },
          },
        },
      },
      null,
      2,
    ),
    'utf8',
  );
  return file;
}

export async function clearVelaLoginConfig(homeDir: string): Promise<void> {
  await rm(join(homeDir, '.amr', 'config.json'), { force: true });
}

function renderFakeVelaScript(options: FakeVelaOptions): string {
  const endpoints = options.endpoints ?? defaultVelaEndpoints();
  return `#!/usr/bin/env node
import { mkdirSync, readFileSync, writeFileSync, existsSync } from 'node:fs';
import { homedir } from 'node:os';
import { dirname, join } from 'node:path';
import { argv, stdin, stdout, stderr, env, exit } from 'node:process';

const ASSISTANT_TEXT = ${JSON.stringify(options.assistantText ?? DEFAULT_ASSISTANT_TEXT)};
const SESSION_ID = env.FAKE_VELA_SESSION_ID || ${JSON.stringify(options.sessionId ?? 'fake-amr-session-1')};
const AUTH_FAIL = ${options.failAuthAtPrompt === true ? 'true' : 'false'};
const AUTH_FAIL_ONCE = ${options.failAuthAtPromptOnce === true ? 'true' : 'false'};
const BALANCE_FAIL = ${options.failBalanceAtPrompt === true ? 'true' : 'false'};
const BALANCE_FAIL_ONCE = ${options.failBalanceAtPromptOnce === true ? 'true' : 'false'};
const MODEL_LIST_INVALID_API_KEY = ${options.failModelListInvalidApiKey === true ? 'true' : 'false'};
const REQUIRE_LOGIN = ${options.requireLoginConfig === false ? 'false' : 'true'};
const REQUIRE_SET_MODEL = ${options.requireSetModel === false ? 'false' : 'true'};

function writeMessage(obj) {
  stdout.write(JSON.stringify(obj) + '\\n');
}
function writeResult(id, result) {
  writeMessage({ jsonrpc: '2.0', id, result });
}
function writeError(id, message) {
  writeMessage({ jsonrpc: '2.0', id, error: { code: -32001, message } });
}
function writeNotification(method, params) {
  writeMessage({ jsonrpc: '2.0', method, params });
}
function currentProfile() {
  return (env.OPEN_DESIGN_AMR_PROFILE || env.VELA_PROFILE || 'local').trim() || 'local';
}
function fakeHomeDir() {
  return homedir();
}
function readLoginConfig() {
  const file = join(fakeHomeDir(), '.amr', 'config.json');
  if (!existsSync(file)) return null;
  try {
    return JSON.parse(readFileSync(file, 'utf8'));
  } catch {
    return null;
  }
}
function hasRuntimeKey() {
  if ((env.VELA_RUNTIME_KEY || '').trim() && (env.VELA_LINK_URL || '').trim()) return true;
  const profile = currentProfile();
  const cfg = readLoginConfig();
  return Boolean(cfg && cfg.profiles && cfg.profiles[profile] && cfg.profiles[profile].runtimeKey);
}
function shouldFailOnce(kind) {
  const file = join(fakeHomeDir(), '.amr', 'fake-vela-state.json');
  let state = {};
  try {
    if (existsSync(file)) state = JSON.parse(readFileSync(file, 'utf8'));
  } catch {
    state = {};
  }
  const key = SESSION_ID + ':' + currentProfile() + ':' + kind;
  if (state[key]) return false;
  state[key] = true;
  mkdirSync(dirname(file), { recursive: true });
  writeFileSync(file, JSON.stringify(state, null, 2), 'utf8');
  return true;
}

if (argv[2] === 'login') {
  const file = join(fakeHomeDir(), '.amr', 'config.json');
  mkdirSync(dirname(file), { recursive: true });
  const profile = currentProfile();
  writeFileSync(file, JSON.stringify({
    profiles: {
      [profile]: {
        runtimeKey: 'fake-runtime-key-0000000000000000000000',
        controlKey: 'fake-control-key-0000000000000000000000',
        apiUrl: ${JSON.stringify(endpoints.apiUrl)},
        linkUrl: ${JSON.stringify(endpoints.linkUrl)},
        user: { id: 'fake-user-id', email: env.FAKE_VELA_LOGIN_USER_EMAIL || 'e2e@example.com', plan: 'free' },
      },
    },
  }, null, 2), 'utf8');
  exit(0);
}

if (argv[2] === 'models') {
  stdout.write('public_model_glm_5    vela\\n');
  exit(0);
}

if (argv[2] === 'model' && argv[3] === 'preset' && argv[4] === '--format' && argv[5] === 'json') {
  stdout.write(${JSON.stringify(PRESET_MODELS_JSON)} + '\\n');
  exit(0);
}

// The daemon asks \`vela model list --all --format json\`; accept it with or without --all, in any flag order.
if (argv[2] === 'model' && argv[3] === 'list' && argv.includes('--format') && argv.includes('json')) {
  if (MODEL_LIST_INVALID_API_KEY) {
    stderr.write('Error: list Link models: API request failed with status 401: invalid_api_key\\n');
    exit(1);
  }
  stdout.write(${JSON.stringify(REMOTE_MODELS_JSON)} + '\\n');
  exit(0);
}

if (REQUIRE_LOGIN && !hasRuntimeKey()) {
  stderr.write('AMR login missing or expired. Please sign in again.\\n');
  exit(1);
}

if (argv[2] === 'billing' && ['preflight', 'funding-check'].includes(argv[3])) {
  const check = ${JSON.stringify(options.fundingCheck ?? null)};
  if (!check) exit(1);
  const workspaceId = argv[argv.indexOf('--workspace-id') + 1];
  const modelId = argv.includes('--model') ? argv[argv.indexOf('--model') + 1] : null;
  const generatedAt = new Date().toISOString();
  stdout.write(JSON.stringify({
    schemaVersion: 1, workspaceId, workspaceMemberId: check.workspaceMemberId,
    modelId, generatedAt, balanceUsd: '0', modelCovered: false,
    funding: 'gateway', fundingCheck: argv[3] === 'funding-check' ? check.status : undefined,
    codingPlan: { workspaceId, generatedAt, eligible: false, tier: null, windows: [] },
  }));
  exit(0);
}

const sessionsWithModel = new Set();
let buffer = '';
stdin.setEncoding('utf8');
stdin.on('data', (chunk) => {
  buffer += chunk;
  const lines = buffer.split('\\n');
  buffer = lines.pop() || '';
  for (const raw of lines) {
    const line = raw.trim();
    if (!line) continue;
    let msg;
    try { msg = JSON.parse(line); } catch { continue; }
    handle(msg);
  }
});
stdin.on('end', () => { stdout.end(); exit(0); });

function handle(msg) {
  const { id, method, params } = msg;
  if (method === 'initialize') {
    writeResult(id, {
      protocolVersion: 1,
      agentCapabilities: { promptCapabilities: { text: true } },
      models: {
        currentModelId: 'glm-5',
        availableModels: [{ modelId: 'glm-5', name: 'glm-5' }],
      },
    });
    return;
  }
  if (method === 'session/new') {
    writeResult(id, {
      sessionId: SESSION_ID,
      models: {
        currentModelId: 'glm-5',
        availableModels: [{ modelId: 'glm-5', name: 'glm-5' }],
      },
    });
    return;
  }
  if (method === 'session/set_model' || method === 'session/set_config_option') {
    const sid = (params && params.sessionId) || SESSION_ID;
    sessionsWithModel.add(sid);
    writeResult(id, {});
    return;
  }
  if (method === 'session/prompt') {
    const sid = (params && params.sessionId) || SESSION_ID;
    if (REQUIRE_SET_MODEL && !sessionsWithModel.has(sid)) {
      writeError(id, 'session/set_model must be called before session/prompt');
      return;
    }
    if (AUTH_FAIL || (AUTH_FAIL_ONCE && shouldFailOnce('auth'))) {
      writeError(id, ${JSON.stringify(PROMPT_AUTH_FAILURE_MESSAGE)});
      return;
    }
    if (BALANCE_FAIL || (BALANCE_FAIL_ONCE && shouldFailOnce('balance'))) {
      writeMessage({
        jsonrpc: '2.0',
        id,
        error: {
          code: -32000,
          message: 'HTTP 429: {"code":"insufficient_balance","message":"insufficient balance"}',
        },
      });
      return;
    }
    writeNotification('session/update', {
      sessionId: sid,
      update: { sessionUpdate: 'agent_message_chunk', content: { type: 'text', text: ASSISTANT_TEXT } },
    });
    writeResult(id, {
      stopReason: 'end_turn',
      usage: { inputTokens: 7, outputTokens: 5, totalTokens: 12 },
    });
    return;
  }
  if (typeof id !== 'undefined') {
    writeMessage({ jsonrpc: '2.0', id, error: { code: -32601, message: 'unknown method ' + method } });
  }
}
`;
}

function defaultVelaEndpoints(): VelaEndpoints {
  return {
    apiUrl: 'http://localhost:18080',
    linkUrl: 'http://localhost:18081',
  };
}