memind/scripts/agent-run-guard.mjs

#!/usr/bin/env node
import fs from 'node:fs';
import path from 'node:path';
import { execFile } from 'node:child_process';
import { promisify } from 'node:util';
import mysql from 'mysql2/promise';

const execFileAsync = promisify(execFile);
const DEFAULT_WORKER_LABEL = 'cn.tkmind.memind-agent-run-worker';
const DEFAULT_PORTAL_LABEL = 'cn.tkmind.memind-portal';

function loadEnvFile(filePath) {
  if (!fs.existsSync(filePath)) return;
  for (const line of fs.readFileSync(filePath, 'utf8').split('\n')) {
    const trimmed = line.trim();
    if (!trimmed || trimmed.startsWith('#')) continue;
    const idx = trimmed.indexOf('=');
    if (idx < 0) continue;
    const key = trimmed.slice(0, idx).trim();
    let value = trimmed.slice(idx + 1).trim();
    if ((value.startsWith('"') && value.endsWith('"')) || (value.startsWith("'") && value.endsWith("'"))) {
      value = value.slice(1, -1);
    }
    if (!process.env[key]) process.env[key] = value;
  }
}

function truthy(value) {
  return ['1', 'true', 'yes', 'on'].includes(String(value ?? '').trim().toLowerCase());
}

function positiveInt(value, fallback) {
  const parsed = Number(value);
  return Number.isFinite(parsed) && parsed > 0 ? Math.floor(parsed) : fallback;
}

function parseArgs(argv) {
  const args = {
    apply: argv.includes('--apply'),
    dryRun: argv.includes('--dry-run') || !argv.includes('--apply'),
    help: argv.includes('--help') || argv.includes('-h'),
  };
  return args;
}

function printHelp() {
  console.log([
    'Usage:',
    '  node scripts/agent-run-guard.mjs [--dry-run]',
    '  node scripts/agent-run-guard.mjs --apply',
    '',
    'Dry-run is the default. --apply can stop the external worker and disable code-run gate in .env.',
  ].join('\n'));
}

function parseMysqlConfig() {
  if (process.env.DATABASE_URL) {
    const url = new URL(process.env.DATABASE_URL);
    if (url.protocol !== 'mysql:') {
      throw new Error(`Unsupported DATABASE_URL scheme for agent run guard: ${url.protocol}`);
    }
    return {
      host: url.hostname,
      port: Number(url.port || 3306),
      user: decodeURIComponent(url.username),
      password: decodeURIComponent(url.password),
      database: url.pathname.replace(/^\/+/, ''),
      charset: 'utf8mb4',
    };
  }
  return {
    host: process.env.MYSQL_HOST,
    port: Number(process.env.MYSQL_PORT || 3306),
    user: process.env.MYSQL_USER,
    password: process.env.MYSQL_PASSWORD,
    database: process.env.MYSQL_DATABASE,
    charset: 'utf8mb4',
  };
}

async function runCommand(command, args) {
  try {
    const { stdout, stderr } = await execFileAsync(command, args, { maxBuffer: 1024 * 1024 });
    return { ok: true, stdout, stderr };
  } catch (err) {
    return {
      ok: false,
      code: err?.code ?? null,
      stdout: err?.stdout ?? '',
      stderr: err?.stderr ?? '',
      message: err instanceof Error ? err.message : String(err),
    };
  }
}

function replaceOrAppendEnv(raw, updates) {
  const pending = new Map(Object.entries(updates));
  const lines = raw.split('\n');
  const next = lines.map((line) => {
    const match = line.match(/^(\s*)([A-Za-z_][A-Za-z0-9_]*)(\s*=).*/);
    if (!match) return line;
    const key = match[2];
    if (!pending.has(key)) return line;
    const value = pending.get(key);
    pending.delete(key);
    return `${key}=${value}`;
  });
  if (pending.size > 0) {
    if (next.length > 0 && next[next.length - 1] !== '') next.push('');
    next.push('## agent-run guard auto-disable');
    for (const [key, value] of pending) next.push(`${key}=${value}`);
  }
  return next.join('\n').replace(/\n*$/, '\n');
}

function timestamp() {
  const d = new Date();
  const pad = (n) => String(n).padStart(2, '0');
  return `${d.getFullYear()}${pad(d.getMonth() + 1)}${pad(d.getDate())}-${pad(d.getHours())}${pad(d.getMinutes())}${pad(d.getSeconds())}`;
}

async function readQueueHealth(now) {
  const conn = await mysql.createConnection(parseMysqlConfig());
  try {
    const [statusRows] = await conn.query(
      `SELECT status, COUNT(*) AS count
       FROM h5_agent_runs
       WHERE status IN ('queued', 'running', 'retryable')
       GROUP BY status`,
    );
    const statusCounts = {};
    for (const row of statusRows) statusCounts[row.status] = Number(row.count ?? 0);

    const [oldestPendingRows] = await conn.query(
      `SELECT MIN(updated_at) AS oldest_updated_at
       FROM h5_agent_runs
       WHERE status IN ('queued', 'retryable')`,
    );
    const oldestPendingUpdatedAt = oldestPendingRows[0]?.oldest_updated_at == null
      ? null
      : Number(oldestPendingRows[0].oldest_updated_at);

    const [oldestRunningRows] = await conn.query(
      `SELECT
         r.id,
         r.request_id,
         r.started_at,
         r.updated_at,
         h.latest_heartbeat_at
       FROM h5_agent_runs r
       LEFT JOIN (
         SELECT run_id, MAX(created_at) AS latest_heartbeat_at
         FROM h5_agent_run_events
         WHERE event_type = 'worker_heartbeat'
         GROUP BY run_id
       ) h ON h.run_id = r.id
       WHERE r.status = 'running'
       ORDER BY COALESCE(h.latest_heartbeat_at, r.started_at) ASC
       LIMIT 1`,
    );
    const oldestRunningStartedAt = oldestRunningRows[0]?.started_at == null
      ? null
      : Number(oldestRunningRows[0].started_at);
    const oldestRunningHeartbeatAt = oldestRunningRows[0]?.latest_heartbeat_at == null
      ? null
      : Number(oldestRunningRows[0].latest_heartbeat_at);
    const oldestRunningHeartbeatAgeMs = oldestRunningRows[0]
      ? Math.max(0, now - Number(oldestRunningRows[0].latest_heartbeat_at ?? oldestRunningRows[0].started_at ?? now))
      : 0;
    const [runningWithoutHeartbeatRows] = await conn.query(
      `SELECT COUNT(*) AS count
       FROM h5_agent_runs r
       LEFT JOIN (
         SELECT run_id, MAX(created_at) AS latest_heartbeat_at
         FROM h5_agent_run_events
         WHERE event_type = 'worker_heartbeat'
         GROUP BY run_id
       ) h ON h.run_id = r.id
       WHERE r.status = 'running' AND h.latest_heartbeat_at IS NULL`,
    );

    const failedWindowMs = positiveInt(process.env.MEMIND_AGENT_RUN_GUARD_FAILED_WINDOW_MS, 10 * 60 * 1000);
    const since = now - failedWindowMs;
    const [failedRows] = await conn.query(
      `SELECT COUNT(*) AS count
       FROM h5_agent_runs
       WHERE status = 'failed' AND updated_at >= ?`,
      [since],
    );
    const [latestFailedRows] = await conn.query(
      `SELECT id, request_id, error_message, updated_at
       FROM h5_agent_runs
       WHERE status = 'failed'
       ORDER BY updated_at DESC
       LIMIT 1`,
    );

    return {
      statusCounts,
      queuedOrRetryable: Number(statusCounts.queued ?? 0) + Number(statusCounts.retryable ?? 0),
      running: Number(statusCounts.running ?? 0),
      oldestPendingUpdatedAt,
      oldestPendingAgeMs: oldestPendingUpdatedAt == null ? 0 : Math.max(0, now - oldestPendingUpdatedAt),
      oldestRunningStartedAt,
      oldestRunningAgeMs: oldestRunningStartedAt == null ? 0 : Math.max(0, now - oldestRunningStartedAt),
      oldestRunningHeartbeatAt,
      oldestRunningHeartbeatAgeMs,
      runningWithoutHeartbeatCount: Number(runningWithoutHeartbeatRows[0]?.count ?? 0),
      latestRunningRun: oldestRunningRows[0] ? {
        id: oldestRunningRows[0].id,
        requestId: oldestRunningRows[0].request_id,
        startedAt: oldestRunningStartedAt,
        updatedAt: oldestRunningRows[0].updated_at == null ? null : Number(oldestRunningRows[0].updated_at),
        heartbeatAt: oldestRunningHeartbeatAt,
        heartbeatAgeMs: oldestRunningHeartbeatAgeMs,
      } : null,
      failedWindowMs,
      failedRecentCount: Number(failedRows[0]?.count ?? 0),
      latestFailedRun: latestFailedRows[0] ? {
        id: latestFailedRows[0].id,
        requestId: latestFailedRows[0].request_id,
        error: latestFailedRows[0].error_message,
        updatedAt: Number(latestFailedRows[0].updated_at ?? 0),
      } : null,
    };
  } finally {
    await conn.end();
  }
}

function evaluateHealth(queue) {
  const thresholds = {
    maxRecentFailures: positiveInt(process.env.MEMIND_AGENT_RUN_GUARD_MAX_RECENT_FAILURES, 3),
    maxPendingAgeMs: positiveInt(process.env.MEMIND_AGENT_RUN_GUARD_MAX_PENDING_AGE_MS, 5 * 60 * 1000),
    maxPendingCount: positiveInt(process.env.MEMIND_AGENT_RUN_GUARD_MAX_PENDING_COUNT, 10),
    maxRunningAgeMs: positiveInt(process.env.MEMIND_AGENT_RUN_GUARD_MAX_RUNNING_AGE_MS, 15 * 60 * 1000),
  };
  const reasons = [];
  if (queue.failedRecentCount >= thresholds.maxRecentFailures) {
    reasons.push(`recent_failed_count ${queue.failedRecentCount} >= ${thresholds.maxRecentFailures}`);
  }
  if (queue.oldestPendingAgeMs >= thresholds.maxPendingAgeMs) {
    reasons.push(`oldest_pending_age_ms ${queue.oldestPendingAgeMs} >= ${thresholds.maxPendingAgeMs}`);
  }
  if (queue.queuedOrRetryable >= thresholds.maxPendingCount) {
    reasons.push(`pending_count ${queue.queuedOrRetryable} >= ${thresholds.maxPendingCount}`);
  }
  if (queue.oldestRunningHeartbeatAgeMs >= thresholds.maxRunningAgeMs) {
    reasons.push(`oldest_running_heartbeat_age_ms ${queue.oldestRunningHeartbeatAgeMs} >= ${thresholds.maxRunningAgeMs}`);
  }
  return { thresholds, reasons, shouldPause: reasons.length > 0 };
}

async function disableCodeRuns({ root, envFile, workerLabel, portalLabel, gui, reasons, dryRun }) {
  const actions = [];
  const updates = {
    MEMIND_AGENT_CODE_RUNS_ENABLED: '0',
    MEMIND_AGENT_RUN_AUTODISPATCH: '1',
  };
  if (dryRun) {
    return {
      applied: false,
      actions: [
        `would backup ${envFile}`,
        `would set ${Object.entries(updates).map(([k, v]) => `${k}=${v}`).join(', ')}`,
        `would stop and disable ${workerLabel}`,
        `would kickstart ${portalLabel}`,
      ],
    };
  }

  const backupRoot = process.env.MEMIND_AGENT_RUN_GUARD_BACKUP_DIR
    || path.join(path.dirname(root), 'memind_backups', `${timestamp()}-agent-run-guard-pause`);
  fs.mkdirSync(backupRoot, { recursive: true });
  const envBackup = path.join(backupRoot, '.env.before');
  fs.copyFileSync(envFile, envBackup);
  actions.push(`backed_up_env:${envBackup}`);

  const raw = fs.readFileSync(envFile, 'utf8');
  const updated = replaceOrAppendEnv(raw, updates);
  const marker = [
    '',
    `# agent-run guard pause at ${new Date().toISOString()}`,
    `# reasons: ${reasons.join('; ')}`,
  ].join('\n');
  fs.writeFileSync(envFile, `${updated.replace(/\n*$/, '\n')}${marker}\n`, 'utf8');
  actions.push(`updated_env:${envFile}`);

  await runCommand('launchctl', ['bootout', `${gui}/${workerLabel}`]);
  await runCommand('launchctl', ['disable', `${gui}/${workerLabel}`]);
  actions.push(`disabled_worker:${workerLabel}`);

  const kick = await runCommand('launchctl', ['kickstart', '-k', `${gui}/${portalLabel}`]);
  actions.push(kick.ok ? `kickstarted_portal:${portalLabel}` : `portal_kickstart_failed:${kick.message}`);
  return { applied: true, backupRoot, actions };
}

const args = parseArgs(process.argv.slice(2));
if (args.help) {
  printHelp();
  process.exit(0);
}

const root = path.join(path.dirname(new URL(import.meta.url).pathname), '..');
const envFile = process.env.MEMIND_ENV_FILE || path.join(root, '.env');
loadEnvFile(envFile);

const enabled = process.env.MEMIND_AGENT_RUN_GUARD_ENABLED == null
  ? true
  : truthy(process.env.MEMIND_AGENT_RUN_GUARD_ENABLED);
const workerLabel = process.env.MEMIND_AGENT_RUN_WORKER_LABEL || DEFAULT_WORKER_LABEL;
const portalLabel = process.env.MEMIND_PORTAL_LABEL || DEFAULT_PORTAL_LABEL;
const gui = `gui/${process.getuid()}`;
const now = Date.now();
const queue = await readQueueHealth(now).catch((err) => ({
  error: err instanceof Error ? (err.message || err.code || err.name) : String(err),
}));
const evaluation = queue.error
  ? { thresholds: {}, reasons: [`queue_health_error ${queue.error}`], shouldPause: false }
  : evaluateHealth(queue);
let pause = { applied: false, actions: [] };

if (enabled && evaluation.shouldPause) {
  pause = await disableCodeRuns({
    root,
    envFile,
    workerLabel,
    portalLabel,
    gui,
    reasons: evaluation.reasons,
    dryRun: args.dryRun,
  });
}

const result = {
  ok: enabled ? !evaluation.shouldPause || pause.applied || args.dryRun : true,
  checkedAt: new Date(now).toISOString(),
  enabled,
  mode: args.apply ? 'apply' : 'dry-run',
  queue,
  evaluation,
  pause,
};

console.log(JSON.stringify(result, null, 2));
process.exit(result.ok ? 0 : 1);