badgr-cli 1.0.31 → 1.0.32

package/src/commands/run.js

@@ -1,13 +1,13 @@
 import { requireApiKey } from '../config.js';
-import { callApi } from '../api.js';
-import { addReceipt, generateReceiptId } from '../store.js';
+import { callApi, terminateDeployment } from '../api.js';
+import { addReceipt, updateReceipt, generateReceiptId } from '../store.js';
+import { normalizeTier, callWithFallback, HIGH_RATE_THRESHOLD } from '../fallback.js';
 
 /**
- * badgr run python train.py --gpu A100           # attached (default)
+ * badgr run python train.py                    # gpu=auto, attached
+ * badgr run python train.py --gpu A100         # specific GPU
  * badgr run --image my/image:latest --gpu L40S --detach
- *
- * Attaches by default: polls status + streams logs until the job finishes,
- * then exits with the job's exit code. Pass --detach to return immediately.
+ * badgr run python train.py --env HF_TOKEN=abc --env DATASET=my/data
  */
 export function parseRunArgs(args) {
   const flags = {};
@@ -18,144 +18,602 @@ export function parseRunArgs(args) {
     if (args[i] === '--image')          { flags.image = args[++i]; i++; continue; }
     if (args[i] === '--count')          { flags.count = parseInt(args[++i], 10); i++; continue; }
     if (args[i] === '--region')         { flags.region = args[++i]; i++; continue; }
+    if (args[i] === '--tier')           { flags.tier = args[++i]; i++; continue; }
     if (args[i] === '--max-price')      { flags.maxPrice = parseFloat(args[++i]); i++; continue; }
     if (args[i] === '--name')           { flags.name = args[++i]; i++; continue; }
-    if (args[i] === '--detach')         { flags.detach = true; i++; continue; }
+    if (args[i] === '--detach')          { flags.detach = true; i++; continue; }
+    if (args[i] === '--fallback')           { flags.fallback = args[++i]; i++; continue; }
+    if (args[i] === '--no-fallback')        { flags.noFallback = true; i++; continue; }
+    if (args[i] === '--strict-capacity')    { flags.noFallback = true; i++; continue; }
+    if (args[i] === '--no-expanded-search') { flags.noFallback = true; i++; continue; }
+    if (args[i] === '--max-runtime')     { flags.maxRuntime = parseFloat(args[++i]); i++; continue; }
+    if (args[i] === '--max-cost')        { flags.maxCost = parseFloat(args[++i]); i++; continue; }
+    if (args[i] === '--env') {
+      const kv = args[++i]; i++;
+      if (!flags.env) flags.env = [];
+      flags.env.push(kv);
+      continue;
+    }
     positional.push(args[i++]);
   }
   return { flags, positional };
 }
 
-// Poll /v1/deployments/:id until terminal status, streaming new log lines.
-async function attachToJob(config, depId, chalk) {
-  const TERMINAL = new Set(['stopped', 'failed', 'completed']);
-  const POLL_MS  = 4000;
-  let seenLines  = 0;
-  let lastStatus = '';
+function parseEnvFlag(envList) {
+  const obj = {};
+  for (const kv of (envList || [])) {
+    const idx = kv.indexOf('=');
+    if (idx > 0) obj[kv.slice(0, idx)] = kv.slice(idx + 1);
+  }
+  return obj;
+}
 
-  while (true) {
-    await new Promise(r => setTimeout(r, POLL_MS));
+// Mirror of backend workload_profile.py — kept in sync for pre-flight display.
+const _PROFILES = {
+  smoke_test:       { label: 'smoke test',            vram: '4 GB',    gpus: ['RTX 3080', 'RTX 3090', 'RTX 4090'] },
+  lora_finetune:    { label: 'fine-tuning (LoRA)',     vram: '40+ GB',  gpus: ['A6000', 'L40S', 'A100'] },
+  image_gen:        { label: 'image generation',       vram: '16+ GB',  gpus: ['RTX 4090', 'A6000', 'L40S'] },
+  inference_small:  { label: 'inference (7B–8B model)', vram: '24+ GB', gpus: ['RTX 4090', 'A6000', 'L40S'] },
+  general:          { label: 'GPU job',                vram: '16+ GB',  gpus: ['RTX 4090', 'RTX 3090', 'A6000', 'L40S'] },
+};
 
-    // Fetch status
-    let dep;
-    try {
-      dep = await callApi(`/deployments/${depId}`, {
-        apiKey: config.apiKey,
-        baseUrl: config.baseUrl,
-      });
-    } catch {
-      // transient network error — keep trying
-      continue;
-    }
+function inferProfileFromCommand(cmdStr) {
+  const s = cmdStr.toLowerCase();
+  if (/print\s*\(|['"]hello/.test(s) && s.length < 100) return 'smoke_test';
+  if (/lora|qlora|finetune|fine[_-]tun|peft/.test(s))    return 'lora_finetune';
+  if (/diffusion|stable.?diff|sdxl|sd.?xl|comfyui|a1111|invoke|kohya/.test(s)) return 'image_gen';
+  if (/vllm|tgi|text.generation.inference/.test(s))       return 'inference_small';
+  if (/\btrain\.py\b/.test(s))                            return 'lora_finetune';
+  return 'general';
+}
 
-    const status = dep.status;
-    if (status !== lastStatus) {
-      if (status === 'running' && lastStatus === 'provisioning') {
-        console.log(chalk.dim('  [running]'));
-      }
-      lastStatus = status;
+function fmtRuntime(ms) {
+  const s = Math.round(ms / 1000);
+  if (s < 60) return `${s}s`;
+  const m = Math.floor(s / 60);
+  return `${m}m ${s % 60}s`;
+}
+
+const HEARTBEAT_WARN_POLLS  = 3;
+const HEARTBEAT_KILL_POLLS  = 15;
+const STARTUP_STATES = new Set(['queued', 'provisioning', 'starting']);
+
+export function classifyFailure(finalStatus, exitCode) {
+  if (finalStatus === 'failed' && (exitCode === null || exitCode === undefined)) return 'infrastructure';
+  if (exitCode !== null && exitCode !== undefined && exitCode !== 0) return 'customer_code';
+  return null;
+}
+
+// Lines the log stream never needs to print — we surface them in the status bar instead.
+const LOG_META_RE = /^\[dep-[^\]]+\] (status|gpu|region|cost|receipt|provider_status|uptime)=/;
+
+// Extract structured values from provider status lines so we can show them nicely.
+function parseProviderLine(line) {
+  const gpuUtil  = line.match(/\bgpu_util=([\d.]+)%/);
+  const cpuUtil  = line.match(/\bcpu_util=([\d.]+)%/);
+  const ssh      = line.match(/\bssh=(\S+)/);
+  const provSt   = line.match(/\bprovider_status=(\S+)/);
+  return {
+    gpuUtil:  gpuUtil  ? parseFloat(gpuUtil[1])  : null,
+    cpuUtil:  cpuUtil  ? parseFloat(cpuUtil[1])  : null,
+    ssh:      ssh      ? ssh[1]                  : null,
+    providerStatus: provSt ? provSt[1]           : null,
+  };
+}
+
+function renderStatusBar(chalk, { elapsedMs, ratePerHour, gpuUtil, cpuUtil, maxRuntimeMs, maxCost }) {
+  const spent = ratePerHour * (elapsedMs / 3_600_000);
+  const parts = [`⏱ ${fmtRuntime(elapsedMs)}`];
+  if (ratePerHour > 0) parts.push(`$${spent.toFixed(4)} spent`);
+  if (gpuUtil !== null) parts.push(`GPU ${gpuUtil.toFixed(0)}%`);
+  if (cpuUtil !== null) parts.push(`CPU ${cpuUtil.toFixed(0)}%`);
+  if (maxRuntimeMs) {
+    const left = Math.max(0, maxRuntimeMs - elapsedMs);
+    parts.push(`${fmtRuntime(left)} left`);
+  }
+  if (maxCost && ratePerHour > 0) {
+    const budgetLeft = Math.max(0, maxCost - spent);
+    parts.push(`$${budgetLeft.toFixed(4)} budget left`);
+  }
+  parts.push('Ctrl+C to stop');
+  return chalk.dim('  ' + parts.join('  •  '));
+}
+
+// Wait for status to leave 'starting'/'queued'/'provisioning'.
+// Returns the dep once it leaves startup states (or the last known state on timeout).
+async function waitForRunning(config, depId, chalk) {
+  const POLL_MS     = 3000;
+  const TIMEOUT_MS  = 5 * 60 * 1000;
+  const startMs     = Date.now();
+  const PHASES = [
+    { afterMs:     0, label: '  Starting container'    },
+    { afterMs: 15000, label: '  Pulling image'         },
+    { afterMs: 60000, label: '  Starting container'    },
+    { afterMs: 180000, label: '  Running command'      },
+  ];
+
+  let lastPhaseIdx = -1;
+
+  const ticker = setInterval(() => {
+    const elapsed = Date.now() - startMs;
+    let phaseIdx = 0;
+    for (let i = 0; i < PHASES.length; i++) {
+      if (elapsed >= PHASES[i].afterMs) phaseIdx = i;
     }
+    if (phaseIdx !== lastPhaseIdx) {
+      process.stdout.write('\r\x1b[2K');
+      process.stdout.write(chalk.dim(PHASES[phaseIdx].label));
+      lastPhaseIdx = phaseIdx;
+    }
+    process.stdout.write('.');
+  }, 1000);
 
-    // Stream any new log lines
-    try {
-      const logData = await callApi(`/deployments/${depId}/logs`, {
+  try {
+    while (Date.now() - startMs < TIMEOUT_MS) {
+      await new Promise(r => setTimeout(r, POLL_MS));
+      const dep = await callApi(`/deployments/${depId}`, {
         apiKey: config.apiKey,
         baseUrl: config.baseUrl,
+        timeoutMs: 10_000,
       });
-      const lines = logData?.logs ?? [];
-      for (let i = seenLines; i < lines.length; i++) {
-        console.log(`  ${chalk.dim(lines[i])}`);
+      if (!STARTUP_STATES.has(dep.status)) {
+        process.stdout.write('\n');
+        return dep;
       }
-      seenLines = lines.length;
-    } catch {
-      // logs not ready yet
     }
+  } finally {
+    clearInterval(ticker);
+    process.stdout.write('\n');
+  }
+
+  return await callApi(`/deployments/${depId}`, {
+    apiKey: config.apiKey,
+    baseUrl: config.baseUrl,
+    timeoutMs: 10_000,
+  });
+}
+
+async function attachToJob(config, depId, { chalk, maxRuntimeMs = null, maxCost = null, ratePerHour = 0, onTeardown, isShuttingDown }) {
+  const TERMINAL = new Set(['stopped', 'failed', 'completed', 'succeeded']);
+  const POLL_MS  = 4000;
+  let seenContent     = new Set();
+  let lastStatus      = '';
+  let consecutiveErrs = 0;
+  let gpuUtil         = null;
+  let cpuUtil         = null;
+  let sshShown        = false;
+  let statusBarActive = false;
+  const startMs       = Date.now();
+
+  // tearing: guards against double-teardown for cap/heartbeat paths within this function.
+  let tearing = false;
+
+  let tickerInterval = null;
+  const startTicker = () => {
+    if (tickerInterval) return;
+    statusBarActive = true;
+    tickerInterval = setInterval(() => {
+      const bar = renderStatusBar(chalk, {
+        elapsedMs: Date.now() - startMs, ratePerHour, gpuUtil, cpuUtil, maxRuntimeMs, maxCost,
+      });
+      process.stdout.write(`\r${bar}   `);
+    }, 1000);
+  };
+  const stopTicker = () => {
+    if (tickerInterval) { clearInterval(tickerInterval); tickerInterval = null; }
+    if (statusBarActive) { process.stdout.write('\r\x1b[2K'); statusBarActive = false; }
+  };
+
+  try {
+    while (true) {
+      // Exit loop if SIGINT handler has started shutdown externally.
+      if (isShuttingDown()) break;
+      if (tearing) break;
+
+      await new Promise(r => setTimeout(r, POLL_MS));
+
+      if (isShuttingDown()) break;
+      if (tearing) break;
+
+      const elapsedMs  = Date.now() - startMs;
+      const spentSoFar = ratePerHour * (elapsedMs / 3_600_000);
+
+      if (maxCost !== null && spentSoFar >= maxCost) {
+        tearing = true;
+        stopTicker();
+        await onTeardown('max-cost');
+        return { status: 'capped', reason: 'max-cost', exitCode: null, runtimeMs: elapsedMs, failureType: null };
+      }
 
-    if (TERMINAL.has(status)) {
-      return status;
+      if (maxRuntimeMs !== null && elapsedMs >= maxRuntimeMs) {
+        tearing = true;
+        stopTicker();
+        await onTeardown('max-runtime');
+        return { status: 'capped', reason: 'max-runtime', exitCode: null, runtimeMs: elapsedMs, failureType: null };
+      }
+
+      let dep;
+      try {
+        dep = await callApi(`/deployments/${depId}`, {
+          apiKey: config.apiKey,
+          baseUrl: config.baseUrl,
+          timeoutMs: 10_000,
+        });
+        consecutiveErrs = 0;
+      } catch {
+        consecutiveErrs++;
+        if (lastStatus === 'running') {
+          const lostSec = Math.round((consecutiveErrs * POLL_MS) / 1000);
+          if (consecutiveErrs === HEARTBEAT_WARN_POLLS) {
+            stopTicker();
+            console.log(chalk.yellow(`\n  ⚠ No heartbeat for ${lostSec}s — cloud machine may be unresponsive`));
+            startTicker();
+          } else if (consecutiveErrs >= HEARTBEAT_KILL_POLLS) {
+            tearing = true;
+            stopTicker();
+            await onTeardown('heartbeat-lost');
+            return { status: 'failed', reason: 'heartbeat-lost', exitCode: null, runtimeMs: elapsedMs, failureType: 'infrastructure' };
+          }
+        }
+        continue;
+      }
+
+      const status = dep.status;
+      if (status !== lastStatus) {
+        if (status === 'running' && lastStatus === 'provisioning') {
+          stopTicker();
+          console.log(chalk.dim('  [running]'));
+        }
+        lastStatus = status;
+      }
+
+      try {
+        const logData = await callApi(`/deployments/${depId}/logs`, {
+          apiKey: config.apiKey,
+          baseUrl: config.baseUrl,
+          timeoutMs: 10_000,
+        });
+        const lines = logData?.logs ?? [];
+
+        for (const line of lines) {
+          const parsed = parseProviderLine(line);
+          if (parsed.gpuUtil  !== null) gpuUtil = parsed.gpuUtil;
+          if (parsed.cpuUtil  !== null) cpuUtil = parsed.cpuUtil;
+
+          if (parsed.ssh && !sshShown) {
+            stopTicker();
+            console.log(`  ${chalk.bold('SSH:')}      ${chalk.cyan(parsed.ssh)}`);
+            sshShown = true;
+            startTicker();
+          }
+
+          if (seenContent.has(line)) continue;
+          seenContent.add(line);
+          if (LOG_META_RE.test(line)) continue;
+          if (/\b(gpu_util|cpu_util|provider_status|uptime)=/.test(line)) continue;
+
+          const isErrorLine = /^error\b/i.test(line) || /Error response from daemon/i.test(line);
+          stopTicker();
+          console.log(`  ${isErrorLine ? chalk.red(line) : chalk.dim(line)}`);
+          startTicker();
+        }
+      } catch {
+        // logs not ready yet
+      }
+
+      if (status === 'running') startTicker();
+
+      if (TERMINAL.has(status)) {
+        stopTicker();
+        const exitCode = dep.exit_code ?? null;
+        return {
+          status,
+          reason: null,
+          exitCode,
+          runtimeMs: Date.now() - startMs,
+          failureType: classifyFailure(status, exitCode),
+        };
+      }
     }
+  } finally {
+    stopTicker();
   }
+
+  // Reached when SIGINT (isShuttingDown) or duplicate tearing flag breaks the loop.
+  // The SIGINT handleShutdown() is managing teardown + exit.
+  return { status: 'interrupted', reason: 'signal', exitCode: null, runtimeMs: Date.now() - startMs, failureType: null };
 }
 
 export async function runCommand(config, args, chalk) {
   const { flags, positional } = parseRunArgs(args);
 
   if (positional.length === 0 && !flags.image) {
-    console.error(chalk.red('Usage: badgr run <command...> --gpu <type>'));
-    console.error(chalk.red('       badgr run --image my/image:latest --gpu A100'));
+    console.error(chalk.red('Usage: badgr run <command...>'));
+    console.error(chalk.red('       badgr run --image my/image:latest'));
     return;
   }
 
   requireApiKey(config);
 
-  const command = positional.length > 0 ? positional : undefined;
-  const gpu     = flags.gpu || 'RTX_4090';
-  const image   = flags.image || (command ? 'python:3.11-slim' : undefined);
-  const detach  = flags.detach || false;
+  // ── Validate flags early ───────────────────────────────────────────────────
+  if (flags.count !== undefined && (!Number.isFinite(flags.count) || flags.count < 1)) {
+    console.error(chalk.red('  ✗ --count must be an integer greater than 0'));
+    process.exitCode = 1;
+    return;
+  }
+  if (flags.maxCost !== undefined && (!Number.isFinite(flags.maxCost) || flags.maxCost <= 0)) {
+    console.error(chalk.red('  ✗ --max-cost must be a number greater than 0'));
+    process.exitCode = 1;
+    return;
+  }
+  if (flags.maxPrice !== undefined && (!Number.isFinite(flags.maxPrice) || flags.maxPrice <= 0)) {
+    console.error(chalk.red('  ✗ --max-price must be a number greater than 0'));
+    process.exitCode = 1;
+    return;
+  }
+  if (flags.region !== undefined && !['US', 'EU', 'AU'].includes(flags.region.toUpperCase())) {
+    console.error(chalk.red('  ✗ --region must be US, EU, or AU'));
+    process.exitCode = 1;
+    return;
+  }
+
+  const command      = positional.length > 0 ? positional : undefined;
+  const cmdStr = command ? command.join(' ') : '';
+  const isSmoke = cmdStr.length < 80 && /print\s*\(|['"]hello/i.test(cmdStr);
+  const inferredImage = isSmoke ? 'python:3.11-alpine' : 'python:3.11-slim';
+  const image = flags.image || (command ? inferredImage : undefined);
+  const detach       = flags.detach || false;
+  const maxRuntimeMs = flags.maxRuntime ? flags.maxRuntime * 60 * 1000 : null;
+  const maxCost      = flags.maxCost ?? null;
+  const envObj       = parseEnvFlag(flags.env);
+  const effectiveTier = normalizeTier(flags.tier);
 
-  console.log(chalk.bold('\n⚡ Running GPU job\n'));
-  if (command) console.log(`  ${chalk.bold('Command:')}  ${command.join(' ')}`);
-  if (image)   console.log(`  ${chalk.bold('Image:')}    ${image}`);
-  console.log(`  ${chalk.bold('GPU:')}      ${gpu}`);
-  if (detach)  console.log(`  ${chalk.dim('(detached — returns immediately)')}`);
-  console.log();
-  console.log(chalk.dim('  Finding best available GPU capacity...'));
+  let gpu = flags.gpu ? flags.gpu.toUpperCase().replace('-', '_') : 'auto';
+
+  if (gpu === 'auto') {
+    console.log(chalk.bold('\n⚡ Running GPU job\n'));
+    if (command) console.log(`  ${chalk.bold('Command:')}  ${command.join(' ')}`);
+    if (image)   console.log(`  ${chalk.bold('Image:')}    ${image}`);
+    if (command) {
+      const profKey = inferProfileFromCommand(cmdStr);
+      const prof = _PROFILES[profKey];
+      console.log();
+      console.log(`  ${chalk.bold('Estimated workload:')}      ${prof.label}`);
+      console.log(`  ${chalk.bold('Estimated minimum VRAM:')}  ${prof.vram}`);
+      if (maxCost)          console.log(`  ${chalk.bold('Max cost:')}                $${maxCost.toFixed(2)}`);
+      if (flags.maxRuntime) console.log(`  ${chalk.bold('Max runtime:')}             ${flags.maxRuntime}min`);
+    }
+    console.log();
+  } else {
+    console.log(chalk.bold('\n⚡ Running GPU job\n'));
+    if (command) console.log(`  ${chalk.bold('Command:')}  ${command.join(' ')}`);
+    if (image)   console.log(`  ${chalk.bold('Image:')}    ${image}`);
+    console.log(`  ${chalk.bold('GPU:')}      ${gpu}`);
+    if (flags.maxRuntime) console.log(`  ${chalk.bold('Max runtime:')} ${flags.maxRuntime}min`);
+    if (maxCost)          console.log(`  ${chalk.bold('Max cost:')}    $${maxCost.toFixed(2)}`);
+    if (flags.maxPrice)   console.log(`  ${chalk.bold('Max price:')}   $${flags.maxPrice.toFixed(2)}/hr`);
+    if (detach)           console.log(`  ${chalk.dim('(detached — returns immediately)')}`);
+    if (flags.env?.length) console.log(`  ${chalk.bold('Env:')}      ${flags.env.join(', ')}`);
+    console.log();
+
+    if (!detach && !flags.maxRuntime && !maxCost) {
+      console.log(chalk.dim('  Tip: add --max-runtime 60 or --max-cost 5.00 to cap spend automatically'));
+    }
+  }
+
+  console.log(chalk.dim('  Finding suitable capacity...'));
+  if (process.env.BADGR_DEBUG === '1' || process.env.BADGR_DEBUG === 'true') {
+    console.log(chalk.dim(`  API: ${config.baseUrl}`));
+  }
+
+  function buildBody(gpuOverride, tierOverride) {
+    const effectiveRegion = flags.region ? flags.region.toUpperCase() : undefined;
+    return {
+      command,
+      image,
+      gpu: (gpuOverride || gpu),
+      gpu_count: flags.count || 1,
+      ...(effectiveRegion ? { region: effectiveRegion } : {}),
+      max_price_per_hour: flags.maxPrice,
+      name: flags.name,
+      tier: tierOverride || effectiveTier,
+      ...(Object.keys(envObj).length > 0 ? { env: envObj } : {}),
+    };
+  }
 
   let dep;
   try {
-    dep = await callApi('/run', {
-      method: 'POST',
-      apiKey: config.apiKey,
-      baseUrl: config.baseUrl,
-      body: {
-        command,
-        image,
-        gpu: gpu.toUpperCase().replace('-', '_'),
-        gpu_count: flags.count || 1,
-        region: flags.region || 'US',
-        max_price_per_hour: flags.maxPrice,
-        name: flags.name,
-      },
-    });
+    dep = await callWithFallback(
+      '/run',
+      { apiKey: config.apiKey, baseUrl: config.baseUrl },
+      (tierOverride) => buildBody(undefined, tierOverride),
+      effectiveTier,
+      chalk,
+      { thing: 'job', cmd: 'badgr run' },
+      { allowTier2Fallback: !flags.noFallback },
+    );
   } catch (err) {
-    console.error(chalk.red(`\n  ✗ Job failed to start: ${err.message}\n`));
-    process.exit(1);
+    if (err.isPaymentRequired) {
+      console.error(chalk.yellow(err.message));
+      const rerun = ['badgr run', ...args].join(' ');
+      console.error(chalk.dim(`After payment, rerun:\n  ${rerun}\n`));
+      process.exitCode = 1;
+      return;
+    }
+    // CapacityError message is pre-formatted with chalk
+    console.error(err.message);
+    process.exitCode = 1;
+    return;
   }
 
   const rcptId = dep.receipt_id || generateReceiptId();
   addReceipt({
-    receiptId:    rcptId,
-    action:       'badgr run',
-    deploymentId: dep.deployment_id,
-    provider:     dep.provider,
-    gpu:          dep.gpu_type,
-    status:       dep.status,
-    createdAt:    new Date().toISOString(),
+    receiptId:     rcptId,
+    action:        'badgr run',
+    deploymentId:  dep.deployment_id,
+    gpu:           dep.gpu_type,
+    providerRoute: dep.provider ?? null,
+    tier:          dep.tier ?? null,
+    maxCost:       maxCost ?? null,
+    maxRuntime:    flags.maxRuntime ?? null,
+    status:        dep.status,
+    createdAt:     new Date().toISOString(),
   });
 
-  console.log(chalk.bold(`\n  Job ID:    ${chalk.cyan(dep.deployment_id)}`));
+  const rate = dep.cost_per_hour || 0;
+
+  console.log(chalk.dim('  Capacity found.\n'));
+  console.log(chalk.bold(`  Job ID:    ${chalk.cyan(dep.deployment_id)}`));
   console.log(`  ${chalk.bold('GPU:')}      ${dep.gpu_type} × ${dep.gpu_count}`);
-  if (dep.cost_per_hour > 0) console.log(`  ${chalk.bold('Rate:')}     $${dep.cost_per_hour.toFixed(2)}/hr`);
+  if (rate > 0) console.log(`  ${chalk.bold('Rate:')}     $${rate.toFixed(2)}/hr`);
+  if (maxCost)  console.log(`  ${chalk.bold('Max cost:')} $${maxCost.toFixed(2)}`);
   console.log(`  ${chalk.bold('Receipt:')}  ${chalk.dim(rcptId)}`);
 
+  if (rate > HIGH_RATE_THRESHOLD && !maxCost) {
+    console.log(chalk.yellow(`\n  Selected capacity rate: $${rate.toFixed(2)}/hr`));
+    console.log(chalk.dim('  Tip: use --max-cost to enforce a hard ceiling.'));
+  }
+
   if (detach) {
     console.log(`\n  ${chalk.bold('Logs:')}     ${dep.logs_url || `badgr logs ${dep.deployment_id}`}`);
     console.log(chalk.dim(`\n  Detached. Track progress: badgr logs ${dep.deployment_id}\n`));
     return;
   }
 
-  // ── Attached mode: stream logs until job completes ───────────────────────
-  console.log(chalk.dim('\n  ── Attaching (Ctrl+C to detach) ──────────────────────────────\n'));
+  // ── Teardown helper ────────────────────────────────────────────────────────
+  // Does NOT call process.exit() — only terminates the deployment and updates receipt.
+  const ratePerHour = dep.cost_per_hour || 0;
+  let teardownCalled = false;
+  let attachStart = Date.now();
+
+  async function teardown(reason) {
+    if (teardownCalled) return;
+    teardownCalled = true;
+
+    const labels = {
+      'max-runtime':    chalk.yellow('\n  ⏱ Max runtime reached — stopping job...'),
+      'max-cost':       chalk.yellow('\n  💰 Spend cap reached — stopping job...'),
+      'heartbeat-lost': chalk.red('\n  ✗ No response from machine — stopping job...'),
+      'interrupted':    chalk.yellow('\n  Stopping job...'),
+    };
+    console.log(labels[reason] ?? chalk.yellow('\n  Stopping job...'));
+
+    try {
+      await terminateDeployment(config, dep.deployment_id);
+    } catch {
+      // terminateDeployment retries 3×; best-effort if all fail
+    }
+
+    const runtimeMs = Date.now() - attachStart;
+    const finalCost = ratePerHour * (runtimeMs / 3_600_000);
+    updateReceipt(rcptId, {
+      status: reason,
+      teardownStatus: 'terminated',
+      runtimeSeconds: Math.round(runtimeMs / 1000),
+      finalCost,
+    });
+    console.log(chalk.dim(`  Runtime: ${fmtRuntime(runtimeMs)}  •  Est. cost: $${finalCost.toFixed(4)}`));
+    console.log(chalk.dim('  Job stopped. Billing ended.\n'));
+  }
+
+  // ── SIGINT handler — installed immediately after we have a deployment ID ───
+  // Covers: queued, provisioning, starting, running phases.
+  // Does NOT call process.exit() inside teardown; exits here after await.
+  let shuttingDown = false;
+
+  async function handleShutdown(reason) {
+    if (shuttingDown) return;
+    shuttingDown = true;
+    try {
+      await teardown(reason);
+    } catch (err) {
+      console.error(chalk.red(`  Warning: teardown may have failed: ${err.message}`));
+    }
+    process.exit(reason === 'interrupted' ? 0 : 1);
+  }
+
+  process.once('SIGINT', () => { void handleShutdown('interrupted'); });
+
+  // ── Wait through startup phases ────────────────────────────────────────────
+  if (STARTUP_STATES.has(dep.status)) {
+    console.log();
+    dep = await waitForRunning(config, dep.deployment_id, chalk);
+  }
+
+  if (dep.status === 'failed') {
+    process.removeListener('SIGINT', handleShutdown);
+    console.error(chalk.red('\n  ✗ Container failed to start (infrastructure error).\n'));
+    console.error(chalk.dim('  The backend retried automatically. All attempts failed.'));
+    console.error(chalk.dim(`  Contact support with receipt ID: ${rcptId}`));
+    console.log();
+    process.exitCode = 1;
+    return;
+  }
+
+  console.log(chalk.dim('\n  ── Running command (Ctrl+C to stop) ────────────────────────────\n'));
+
+  attachStart = Date.now();
+  const { status: finalStatus, exitCode, runtimeMs, failureType } = await attachToJob(config, dep.deployment_id, {
+    chalk,
+    maxRuntimeMs,
+    maxCost,
+    ratePerHour,
+    onTeardown: teardown,
+    isShuttingDown: () => shuttingDown,
+  });
+
+  // Remove SIGINT handler — job is done (or SIGINT was handled)
+  process.removeListener('SIGINT', handleShutdown);
 
-  const finalStatus = await attachToJob(config, dep.deployment_id, chalk);
+  // 'interrupted' = SIGINT handler is managing teardown + exit — don't duplicate
+  if (finalStatus === 'interrupted') return;
 
   console.log();
 
-  if (finalStatus === 'failed') {
-    console.error(chalk.red(`\n  ✗ Job failed (${dep.deployment_id})\n`));
-    console.error(chalk.dim(`  Logs: badgr logs ${dep.deployment_id}\n`));
-    process.exit(1);
-  } else {
-    console.log(chalk.green(`\n  ✓ Job complete (${finalStatus})\n`));
+  const finalCost = ratePerHour * (runtimeMs / 3_600_000);
+  updateReceipt(rcptId, {
+    status:        finalStatus,
+    exitCode,
+    runtimeSeconds: Math.round(runtimeMs / 1000),
+    finalCost,
+    failureType,
+    teardownStatus: (finalStatus === 'completed' || finalStatus === 'succeeded') ? 'terminated' : 'failed',
+  });
+
+  console.log(`  ${chalk.bold('Runtime:')}  ${fmtRuntime(runtimeMs)}`);
+  if (ratePerHour > 0) {
+    console.log(`  ${chalk.bold('Cost:')}     $${finalCost.toFixed(4)}  (${chalk.dim(`$${ratePerHour.toFixed(2)}/hr`)})`);
+  }
+  if (exitCode !== null && exitCode !== undefined) {
+    console.log(`  ${chalk.bold('Exit code:')} ${exitCode !== 0 ? chalk.red(exitCode) : chalk.green(exitCode)}`);
+  }
+  console.log(`  ${chalk.bold('Receipt:')}  ${chalk.dim(rcptId)}`);
+
+  // 'capped' = max-runtime or max-cost path; teardown message already printed
+  if (finalStatus === 'capped') {
+    process.exitCode = 1;
+    return;
+  }
+
+  if (finalStatus === 'failed' || (exitCode !== null && exitCode !== 0)) {
+    if (failureType === 'infrastructure') {
+      console.error(chalk.red(`\n  ✗ Machine failure — this is not your code.\n`));
+      console.error(chalk.dim('  Contact support with your receipt ID for a refund.'));
+    } else {
+      console.error(chalk.red(`\n  ✗ Job failed (exit ${exitCode ?? 'unknown'})\n`));
+      console.error(chalk.dim(`  Check logs: badgr logs ${dep.deployment_id}`));
+    }
+    console.log();
+    process.exitCode = exitCode ?? 1;
+    return;
+  }
+
+  if ((finalStatus === 'completed' || finalStatus === 'succeeded') && (exitCode === 0 || exitCode === null)) {
+    try {
+      await terminateDeployment(config, dep.deployment_id);
+    } catch { /* already stopped */ }
+    console.log(chalk.green(`\n  ✓ Complete`));
+    console.log(chalk.dim(`  Billing ended`));
+    console.log();
   }
 }