cjyyzx/gateway/scripts/bootstrap-lobechat-users.js

#!/usr/bin/env node
/* eslint-disable no-console */
/**
 * Bootstrap every existing LobeChat user so the UI is immediately usable:
 *   1. Enables the built-in "azure" (Azure OpenAI) provider with encrypted keyVaults.
 *   2. Inserts one "azure/<deploymentName>" ai_models row enabled=true.
 *   3. Installs one customPlugin "enterprise-gateway" pointing at the gateway
 *      manifest URL (identity-aware; each user sees only their own allowed tools).
 *
 * Runs inside the lobechat container so it can use the bundled `pg` module and
 * reach `db:5432` / `gateway:3001` on the compose network. Idempotent:
 * re-running does not duplicate or error.
 *
 * Required env (passed through by the wrapper script):
 *   LOBECHAT_DATABASE_URL   postgres URL for the lobechat DB
 *   KEY_VAULTS_SECRET       same secret LobeChat uses to decrypt key_vaults
 *   AZURE_OPENAI_API_KEY    Azure OpenAI key
 *   AZURE_OPENAI_ENDPOINT   https://<resource>.openai.azure.com (or .cognitiveservices.azure.com)
 *   AZURE_OPENAI_API_VERSION e.g. 2025-04-01-preview
 *   AZURE_OPENAI_DEPLOYMENT deployment name (used as the model id)
 *   GATEWAY_MANIFEST_URL    defaults to http://gateway:3001/api/lobechat/manifest
 */

const crypto = require('crypto');
const { Client } = require('pg');

const DB_URL = process.env.LOBECHAT_DATABASE_URL;
const KEY_VAULTS_SECRET = process.env.KEY_VAULTS_SECRET;
const AZ_KEY = process.env.AZURE_OPENAI_API_KEY;
const AZ_ENDPOINT = process.env.AZURE_OPENAI_ENDPOINT;
const AZ_VERSION = process.env.AZURE_OPENAI_API_VERSION || '2025-04-01-preview';
const AZ_DEPLOY = process.env.AZURE_OPENAI_DEPLOYMENT || 'gpt-5.4';
const MANIFEST_URL =
  process.env.GATEWAY_MANIFEST_URL || 'http://gateway:3001/api/lobechat/manifest';

function die(msg) {
  console.error('[bootstrap] FATAL:', msg);
  process.exit(1);
}

if (!DB_URL) die('LOBECHAT_DATABASE_URL is required');
if (!KEY_VAULTS_SECRET) die('KEY_VAULTS_SECRET is required');
if (!AZ_KEY || !AZ_ENDPOINT) die('AZURE_OPENAI_API_KEY and AZURE_OPENAI_ENDPOINT are required');

/**
 * Reproduces src/server/modules/KeyVaultsEncrypt. LobeChat stores
 * `${ivHex}:${authTagHex}:${cipherHex}` using AES-GCM with the raw bytes of
 * KEY_VAULTS_SECRET base64-decoded (must be 16/24/32 bytes).
 */
function encryptKeyVaults(plaintext) {
  const rawKey = Buffer.from(KEY_VAULTS_SECRET, 'base64');
  if (![16, 24, 32].includes(rawKey.length)) {
    die(
      `KEY_VAULTS_SECRET must decode to 16/24/32 bytes; got ${rawKey.length}. ` +
        `Regenerate with: openssl rand -base64 32 (and update .env + restart lobechat).`,
    );
  }
  const iv = crypto.randomBytes(12);
  const cipher = crypto.createCipheriv('aes-256-gcm', rawKey, iv);
  const enc = Buffer.concat([cipher.update(plaintext, 'utf8'), cipher.final()]);
  const authTag = cipher.getAuthTag();
  return `${iv.toString('hex')}:${authTag.toString('hex')}:${enc.toString('hex')}`;
}

async function main() {
  const client = new Client({ connectionString: DB_URL });
  await client.connect();
  console.log('[bootstrap] connected to', DB_URL.replace(/:[^:@]*@/, ':***@'));

  const { rows: users } = await client.query('SELECT id, email FROM users ORDER BY created_at');
  console.log(`[bootstrap] found ${users.length} user(s)`);

  // keyVaults shape for Azure OpenAI — see packages/types/src/user/settings/keyVaults.ts
  const keyVaultsJson = JSON.stringify({
    apiKey: AZ_KEY,
    apiVersion: AZ_VERSION,
    baseURL: AZ_ENDPOINT,
    endpoint: AZ_ENDPOINT, // deprecated alias; included for older runtime code paths
  });

  // CustomPlugin manifest + params. LobeChat fetches `manifestUrl` on first use
  // and caches into the `manifest` column; we pre-populate nothing so the
  // identity-aware manifest is always re-fetched per user.
  const pluginIdentifier = 'enterprise-gateway';
  const customParams = {
    apiMode: 'simple',
    avatar: '🏢',
    description:
      'Identity-aware Enterprise Gateway tools. Exposed tool list filtered by caller RBAC.',
    manifestMode: 'url',
    manifestUrl: MANIFEST_URL,
  };

  let providersInserted = 0;
  let modelsInserted = 0;
  let pluginsInserted = 0;

  for (const u of users) {
    // Fresh ciphertext per user so two users sharing a leaked row hash nothing.
    const encKV = encryptKeyVaults(keyVaultsJson);

    // --- ai_providers (builtin "azure") -----------------------------------
    const prov = await client.query(
      `INSERT INTO ai_providers (id, user_id, enabled, key_vaults, source, name, logo, settings, config)
       VALUES ($1, $2, TRUE, $3, 'builtin', 'Azure OpenAI', NULL, '{}'::jsonb, '{}'::jsonb)
       ON CONFLICT (id, user_id) DO UPDATE
         SET enabled = TRUE,
             key_vaults = EXCLUDED.key_vaults
         WHERE ai_providers.key_vaults IS NULL OR ai_providers.key_vaults = ''
       RETURNING id`,
      ['azure', u.id, encKV],
    );
    if (prov.rowCount) providersInserted++;

    // --- ai_models (one row for the deployment) ---------------------------
    const mdl = await client.query(
      `INSERT INTO ai_models
         (id, provider_id, user_id, enabled, display_name, type, source, config, abilities, parameters, settings)
       VALUES ($1, 'azure', $2, TRUE, $3, 'chat', 'custom',
               $4::jsonb,
               '{"functionCall":true,"vision":true}'::jsonb,
               '{}'::jsonb, '{}'::jsonb)
       ON CONFLICT (id, provider_id, user_id) DO UPDATE SET enabled = TRUE
       RETURNING id`,
      [
        AZ_DEPLOY,
        u.id,
        `${AZ_DEPLOY} (Azure)`,
        JSON.stringify({ deploymentName: AZ_DEPLOY }),
      ],
    );
    if (mdl.rowCount) modelsInserted++;

    // --- user_installed_plugins (customPlugin) ----------------------------
    const plg = await client.query(
      `INSERT INTO user_installed_plugins
         (user_id, identifier, type, manifest, settings, custom_params, source)
       VALUES ($1, $2, 'customPlugin', NULL, '{}'::jsonb, $3::jsonb, 'custom')
       ON CONFLICT (user_id, identifier) DO UPDATE
         SET custom_params = EXCLUDED.custom_params
       RETURNING identifier`,
      [u.id, pluginIdentifier, JSON.stringify(customParams)],
    );
    if (plg.rowCount) pluginsInserted++;

    console.log(`[bootstrap] ✓ ${u.email || u.id}`);
  }

  console.log(
    `[bootstrap] done. providers=${providersInserted} models=${modelsInserted} plugins=${pluginsInserted}`,
  );
  await client.end();
}

main().catch((err) => {
  console.error('[bootstrap] error:', err);
  process.exit(1);
});