clawdie-ai/setup/agent-jails.ts

/**
 * setup/agent-jails.ts — Provision per-specialist agent jails.
 *
 * Creates db-worker, git-worker, ctrl-worker jails from infra/jails.yaml.
 * Each gets pi + aider installed, scoped secrets, and nullfs mounts for
 * its domain. PF rules restrict cross-jail network access.
 *
 * Run: just setup-agent-jails  (or npm run install -- --from agent-jails)
 */
import fs from 'fs';
import path from 'path';

import {
  AGENT_NAME,
  SUBNET_BASE,
  DB_HOST,
  PROJECT_ROOT,
} from '../src/config.js';
import { logger } from '../src/logger.js';
import { readEnvFile } from '../src/env.js';
import { loadPackageList, mountPkgCacheInJail } from './packages.js';
import { commandExists, getPlatform, isRoot } from './platform.js';
import { emitStatus } from './status.js';
import {
  bastille,
  jailExists,
  detectFreeBSDRelease,
  jailRoot,
  resolveJailName,
} from './bastille-helpers.js';
import {
  loadJailRegistry,
  resolveJailIp,
  getAgentJailDef,
  parseMountSpec,
} from '../src/jail-schema.js';

const LOG = 'logs/setup.log';

const SHARED_LLM_KEYS = [
  'ANTHROPIC_API_KEY',
  'OPENAI_API_KEY',
  'OPENROUTER_API_KEY',
  'ZAI_API_KEY',
  'GEMINI_API_KEY',
];

const DOMAIN_KEY_ENV_MAP: Record<string, string[]> = {
  'db-admin': [
    'DB_HOST',
    'DB_PORT',
    'SKILLS_DB_URL',
    'MEMORY_DB_URL',
    'OPS_DB_URL',
    'SKILLS_DB_USER',
    'SKILLS_DB_PASSWORD',
    'MEMORY_DB_USER',
    'MEMORY_DB_PASSWORD',
    'OPS_DB_USER',
    'OPS_DB_PASSWORD',
  ],
  'git-admin': [
    'REMOTE_GIT_URL',
    'GIT_MIRROR_URLS',
    'GIT_SSH_KEY_PATH',
    'FORGEJO_TOKEN',
  ],
  coordinator: ['CONTROLPLANE_API_KEY'],
};

const EXTRA_PACKAGES: Record<string, string[]> = {
  'db-admin': ['postgresql18-client'],
  'git-admin': [],
  coordinator: [],
};

function ensurePiInstalled(jailName: string): void {
  const exists = bastille('cmd', jailName, '/bin/sh', '-lc', 'command -v pi');
  if (exists.ok) return;

  logger.info({ jailName }, 'Installing pi inside agent jail');
  const install = bastille(
    'cmd',
    jailName,
    'npm',
    'install',
    '-g',
    '@mariozechner/pi-coding-agent',
  );
  if (!install.ok) {
    throw new Error(`pi install failed in ${jailName}: ${install.output}`);
  }
}

function buildJailEnv(
  specialist: string,
  allEnv: Record<string, string | undefined>,
): string {
  const lines: string[] = [];

  const sharedKey =
    allEnv['ANTHROPIC_API_KEY'] ||
    allEnv['OPENAI_API_KEY'] ||
    allEnv['OPENROUTER_API_KEY'] ||
    allEnv['ZAI_API_KEY'] ||
    allEnv['GEMINI_API_KEY'];

  if (sharedKey) {
    lines.push(`ANTHROPIC_API_KEY=${sharedKey}`);
  }

  const domainKeys = DOMAIN_KEY_ENV_MAP[specialist] ?? [];
  for (const key of domainKeys) {
    const val = allEnv[key];
    if (val) lines.push(`${key}=${val}`);
  }

  return lines.join('\n') + '\n';
}

export async function run(args: string[]): Promise<void> {
  if (getPlatform() !== 'freebsd') {
    emitStatus('SETUP_AGENT_JAILS', {
      STATUS: 'failed',
      ERROR: 'unsupported_platform',
      LOG,
    });
    process.exit(1);
  }
  if (!isRoot()) {
    emitStatus('SETUP_AGENT_JAILS', {
      STATUS: 'failed',
      ERROR: 'requires_root',
      LOG,
    });
    throw new Error('setup_agent_jails_requires_root');
  }
  if (!commandExists('bastille')) {
    emitStatus('SETUP_AGENT_JAILS', {
      STATUS: 'failed',
      ERROR: 'missing_bastille',
      LOG,
    });
    throw new Error('missing_bastille');
  }

  const registry = loadJailRegistry();
  const gateway = process.env.WARDEN_GATEWAY || `${SUBNET_BASE}.1`;
  const bridge = process.env.WARDEN_BRIDGE || registry.bridge;
  const release = detectFreeBSDRelease();
  const internalDomain =
    process.env.AGENT_INTERNAL_DOMAIN || `${AGENT_NAME}.home.arpa`;
  const allEnv = readEnvFile([
    ...SHARED_LLM_KEYS,
    ...Object.values(DOMAIN_KEY_ENV_MAP).flat(),
  ]);

  const specialists = ['db-admin', 'git-admin', 'coordinator'];
  const results: Record<
    string,
    { jailName: string; ip: string; created: boolean }
  > = {};

  for (const specialist of specialists) {
    const entry = getAgentJailDef(registry, specialist);
    if (!entry) {
      logger.warn({ specialist }, 'No agent jail definition found in registry');
      continue;
    }

    const { role, def } = entry;
    const jailName = resolveJailName({
      agentName: AGENT_NAME,
      role,
      legacyNames: [role],
    });
    const ip = resolveJailIp(registry, role);
    const hostname = `${role}.${internalDomain}`;

    const exists = jailExists(jailName);

    if (!exists) {
      if (args.includes('--destroy') || args.includes('--recreate')) {
        logger.info(
          { jailName, specialist },
          'Skipping creation (--destroy/--recreate passed but jail does not exist)',
        );
      }
      logger.info(
        { jailName, ip, release, specialist },
        `Creating ${role} jail`,
      );
      const create = bastille(
        'create',
        ...(def.thick ? ['-T'] : []),
        ...(def.vnet ? ['-B'] : []),
        '-g',
        gateway,
        jailName,
        release,
        `${ip}/24`,
        bridge,
      );
      if (!create.ok) {
        throw new Error(`bastille create failed for ${role}: ${create.output}`);
      }
      bastille('config', jailName, 'set', 'host.hostname', hostname);
      bastille('restart', jailName);
    } else {
      logger.info(
        { jailName, specialist },
        `${role} jail already exists, skipping creation`,
      );
    }

    mountPkgCacheInJail(jailName);

    const pkgs = loadPackageList('agent-worker-jail.txt');
    const extra = EXTRA_PACKAGES[specialist] ?? [];
    const allPkgs = [...pkgs, ...extra];

    if (allPkgs.length > 0) {
      const pkg = bastille('pkg', jailName, 'install', '-y', ...allPkgs);
      if (!pkg.ok) {
        logger.warn(
          { output: pkg.output, specialist },
          `${role} jail package install had warnings`,
        );
      }
    }

    ensurePiInstalled(jailName);

    const chsh = bastille(
      'cmd',
      jailName,
      'chsh',
      '-s',
      '/usr/local/bin/bash',
      'root',
    );
    if (!chsh.ok) {
      logger.warn({ output: chsh.output }, `chsh failed in ${role}`);
    }

    const envContent = buildJailEnv(specialist, allEnv);
    const jailEnvPath = path.join(jailRoot(jailName), 'root', '.env.agent');
    fs.mkdirSync(path.dirname(jailEnvPath), { recursive: true });
    fs.writeFileSync(jailEnvPath, envContent);
    logger.info(
      { jailName, specialist, envPath: jailEnvPath },
      'Wrote scoped .env.agent',
    );

    if (entry.agentJail.mounts && entry.agentJail.mounts.length > 0) {
      for (const mountSpec of entry.agentJail.mounts) {
        const mount = parseMountSpec(mountSpec);
        const hostDir = resolveMountHostPath(mount.hostToken, allEnv);
        if (!hostDir || !fs.existsSync(hostDir)) {
          logger.warn(
            { specialist, mountHost: mount.hostToken, resolvedPath: hostDir },
            `Skipping mount — host path not found`,
          );
          continue;
        }

        const fstabPath = `/usr/local/bastille/jails/${jailName}/fstab`;

        if (fs.existsSync(fstabPath)) {
          const existing = fs.readFileSync(fstabPath, 'utf-8');
          if (existing.includes(mount.jailPath)) {
            logger.debug({ specialist, mountPoint: mount.jailPath }, 'Mount already in fstab');
            continue;
          }
        }

        bastille('mount', jailName, hostDir, mount.jailPath, 'nullfs', mount.mode);
        logger.info(
          { specialist, hostDir, jailPoint: mount.jailPath, mode: mount.mode },
          'Mounted directory into agent jail',
        );
      }
    }

    results[specialist] = { jailName, ip, created: !exists };
  }

  emitStatus('SETUP_AGENT_JAILS', {
    STATUS: 'success',
    JAILS: JSON.stringify(results),
    LOG,
  });
}

function resolveMountHostPath(
  token: string,
  env: Record<string, string | undefined>,
): string | null {
  if (token === 'MIGRATIONS_DIR') {
    return path.join(PROJECT_ROOT, 'docs', 'internal', 'sql');
  }
  if (token === 'SKILLS_DIR') {
    return path.join(PROJECT_ROOT, '.agent', 'skills');
  }
  if (token === 'IDENTITIES_DIR') {
    return path.join(PROJECT_ROOT, '.agent', 'identities');
  }
  if (token === 'REPO_DIR') {
    const repoRoot = env['GIT_STORAGE_ROOT'];
    if (repoRoot) return repoRoot;
    return path.join(PROJECT_ROOT);
  }
  const envVal = env[token];
  if (envVal && fs.existsSync(envVal)) return envVal;
  return null;
}