badgr-cli 1.0.44 → 1.0.46

package/src/commands/run.js

@@ -8,7 +8,8 @@ import { addReceipt, updateReceipt, generateReceiptId } from '../store.js';
 import { normalizeTier, callWithFallback, HIGH_RATE_THRESHOLD } from '../fallback.js';
 import { formatCliError } from '../errors.js';
 import { TEMPLATE_MAP, buildTemplateFlags, parseTemplateOverrides } from '../catalog.js';
-import { stage as _stage, stageDone as _stageDone, writeBlock as _writeBlock, clearBlock as _clearBlock, renderLiveBlock as _renderLiveBlock } from '../progress.js';
+import { stage as _stage, stageDone as _stageDone, writeBlock as _writeBlock, clearBlock as _clearBlock, renderLiveBlock as _renderLiveBlock, printFailureClass as _printFailureClass } from '../progress.js';
+import { detectWorkload, workloadTypeLabel } from '../detect.js';
 
 /**
  * Flow 1 — local project (primary):
@@ -56,6 +57,10 @@ export function parseRunArgs(args) {
     if (flagArgs[i] === '--save')            { flags.save = flagArgs[++i]; i++; continue; }
     if (flagArgs[i] === '--workspace')       { flags.workspace = flagArgs[++i]; i++; continue; }
     if (flagArgs[i] === '--cmd')             { flags.cmd = flagArgs[++i]; i++; continue; }
+    if (flagArgs[i] === '--output')          { flags.output = flagArgs[++i]; i++; continue; }
+    if (flagArgs[i] === '--checkpoint')      { flags.checkpoint = flagArgs[++i]; i++; continue; }
+    if (flagArgs[i] === '--retry-safe')      { flags.retrySafe = true; i++; continue; }
+    if (flagArgs[i] === '--resume-cmd')      { flags.resumeCmd = flagArgs[++i]; i++; continue; }
     if (flagArgs[i] === '--env') {
       const kv = flagArgs[++i]; i++;
       if (!flags.env) flags.env = [];
@@ -339,6 +344,8 @@ async function attachToJob(config, depId, { chalk, maxRuntimeMs = null, maxCost
           exitCode,
           runtimeMs: Date.now() - startMs,
           failureType: classifyFailure(status, exitCode),
+          failureClass: dep.failure_class ?? null,
+          nextAction: dep.next_action ?? null,
         };
       }
     }
@@ -357,6 +364,7 @@ const _KNOWN_RUN_FLAGS = new Set([
   '--detach', '--fallback', '--no-fallback', '--strict-capacity',
   '--no-expanded-search', '--max-runtime', '--max-cost', '--min-vram', '--env',
   '--dry-run', '--cmd', '--save', '--workspace',
+  '--output', '--checkpoint', '--retry-safe', '--resume-cmd',
 ]);
 
 // Directories and files always excluded from project zip uploads.
@@ -497,10 +505,43 @@ export async function runCommand(config, args, chalk) {
   const isGitHubUrl  = firstArg && _isGitHubUrl(firstArg);
   const isCodeSource = isLocalPath || isGitHubUrl;
 
+  // Local paths can be inspected on disk, so a missing --cmd doesn't have to be
+  // an error — try to infer it (and output/checkpoint conventions) the same way
+  // `badgr detect .` would, before ever provisioning anything.
+  let detectionReport = null;
+  if (isLocalPath && !flags.cmd) {
+    detectionReport = detectWorkload(firstArg);
+    if (detectionReport.command && detectionReport.confidence !== 'low') {
+      flags.cmd = detectionReport.command;
+      console.log(chalk.dim(`\n  Detected command (${detectionReport.confidence} confidence): ${detectionReport.command}`));
+      console.log(chalk.dim(`  Workload: ${workloadTypeLabel(detectionReport.workloadType)}  •  badgr detect . for the full report  •  --cmd to override`));
+    }
+    if (!flags.output && detectionReport.outputs?.length) flags.output = detectionReport.outputs[0];
+    if (!flags.checkpoint && detectionReport.checkpoints?.length) flags.checkpoint = detectionReport.checkpoints[0];
+  }
+
+  // Detection couldn't confidently name a command — ask for just that one
+  // missing essential (interactively, if we have a terminal to ask on).
+  if (isLocalPath && !flags.cmd && process.stdin.isTTY && process.stdout.isTTY) {
+    console.log(chalk.dim(`\n  Badgr couldn't confidently detect a command to run in ${path.resolve(firstArg)}.`));
+    try {
+      const { input } = await import('@inquirer/prompts');
+      const answer = await input({
+        message: 'Which command should run?',
+        default: detectionReport?.command || undefined,
+        validate: v => v.trim() ? true : 'A command is required',
+      });
+      flags.cmd = answer.trim();
+    } catch {
+      // Ctrl+C or a non-interactive stdin that lied about isTTY — fall through to the hard error below.
+    }
+  }
+
   if (isCodeSource && !flags.cmd) {
     console.error(chalk.red(`\n  ✗ --cmd is required when running from a ${isLocalPath ? 'local path' : 'GitHub URL'}.\n`));
     if (isLocalPath) {
       console.error(chalk.dim('  Example: badgr run . --cmd "python train.py" --max-cost 5\n'));
+      console.error(chalk.dim('  Or check what Badgr detects first: badgr detect .\n'));
     } else {
       console.error(chalk.dim('  Example: badgr run https://github.com/user/repo --cmd "python train.py" --max-cost 5\n'));
     }
@@ -550,6 +591,20 @@ export async function runCommand(config, args, chalk) {
     return;
   }
 
+  if (!flags.maxCost && !flags.dryRun && isLocalPath && process.stdin.isTTY && process.stdout.isTTY) {
+    try {
+      const { input } = await import('@inquirer/prompts');
+      const answer = await input({
+        message: 'What max cost ($) should cap this run?',
+        default: '5',
+        validate: v => (Number.isFinite(parseFloat(v)) && parseFloat(v) > 0) ? true : 'Enter a number greater than 0',
+      });
+      flags.maxCost = parseFloat(answer);
+    } catch {
+      // fall through to the hard error below
+    }
+  }
+
   if (!flags.maxCost && !flags.dryRun) {
     console.error(chalk.red('\n  ✗ --max-cost is required for run workloads.\n'));
     console.error(chalk.dim('  Example: badgr run --gpu RTX_4090 --image node:20 --max-cost 5 -- node script.js\n'));
@@ -578,7 +633,16 @@ export async function runCommand(config, args, chalk) {
   const maxRuntimeMs = effectiveMaxRuntime * 60 * 1000;
 
   const maxCost      = flags.maxCost ?? null;
-  const envObj       = parseEnvFlag(flags.env);
+  // BADGR_OUTPUT_DIR / BADGR_CHECKPOINT_DIR / BADGR_RETRY_SAFE are a
+  // convention for custom.run containers to write recoverable state outside
+  // the pod — Badgr doesn't enforce what the container does with them, it
+  // just wires the directories through and surfaces --resume-cmd on failure.
+  // Explicit --env always wins if the user passes the same key directly.
+  const conventionEnv = {};
+  if (flags.output)     conventionEnv.BADGR_OUTPUT_DIR = flags.output;
+  if (flags.checkpoint) conventionEnv.BADGR_CHECKPOINT_DIR = flags.checkpoint;
+  if (flags.retrySafe)  conventionEnv.BADGR_RETRY_SAFE = '1';
+  const envObj       = { ...conventionEnv, ...parseEnvFlag(flags.env) };
   const effectiveTier = normalizeTier(flags.tier);
 
   const gpu = flags.gpu ? flags.gpu.toUpperCase().replace('-', '_') : undefined;
@@ -595,6 +659,10 @@ export async function runCommand(config, args, chalk) {
     if (maxCost)          console.log(`  ${chalk.bold('Max cost:')}     $${maxCost}`);
     if (flags.maxPrice)   console.log(`  ${chalk.bold('Max price:')}    $${flags.maxPrice}/hr`);
     if (flags.env?.length) console.log(`  ${chalk.bold('Env:')}         ${flags.env.join(', ')}`);
+    if (flags.output)     console.log(`  ${chalk.bold('Output:')}       ${flags.output}`);
+    if (flags.checkpoint) console.log(`  ${chalk.bold('Checkpoint:')}   ${flags.checkpoint}`);
+    if (flags.retrySafe)  console.log(`  ${chalk.bold('Retry-safe:')}   enabled`);
+    if (flags.resumeCmd)  console.log(`  ${chalk.bold('Resume cmd:')}   ${flags.resumeCmd}`);
     console.log(chalk.dim('\n  Remove --dry-run to provision.\n'));
     return;
   }
@@ -616,6 +684,9 @@ export async function runCommand(config, args, chalk) {
   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(', ')}`);
+  if (flags.output)     console.log(`  ${chalk.bold('Output:')}     ${flags.output}`);
+  if (flags.checkpoint) console.log(`  ${chalk.bold('Checkpoint:')} ${flags.checkpoint}`);
+  if (flags.retrySafe)  console.log(`  ${chalk.bold('Retry-safe:')} enabled`);
   console.log();
 
 
@@ -775,6 +846,7 @@ export async function runCommand(config, args, chalk) {
     console.log(chalk.dim(_stage(STAGE_TOTAL, STAGE_TOTAL, stageLabels[reason] ?? 'Stopped')));
     console.log(chalk.dim(`  Runtime: ${fmtRuntime(runtimeMs)}  •  Estimated cost: ~$${finalCost.toFixed(2)}`));
     _printFinalInfo(chalk, { exitCode: null, teardownOk, jobId: dep.deployment_id, rcptId });
+    if (flags.resumeCmd) console.log(`  ${chalk.bold('Resume:')}    ${flags.resumeCmd}`);
     console.log();
   }
 
@@ -805,6 +877,7 @@ export async function runCommand(config, args, chalk) {
   if (dep.status === 'failed') {
     process.removeListener('SIGINT', handleShutdown);
     console.error(formatCliError('JOB_INFRASTRUCTURE_FAILURE', { receiptId: rcptId }, chalk));
+    _printFailureClass(chalk, dep);
     process.exitCode = 1;
     return;
   }
@@ -813,7 +886,7 @@ export async function runCommand(config, args, chalk) {
   console.log(chalk.dim(`\n  [${stageN}/${STAGE_TOTAL}] Running command (Ctrl+C to stop)`));
 
   attachStart = Date.now();
-  const { status: finalStatus, exitCode, runtimeMs, failureType } = await attachToJob(config, dep.deployment_id, {
+  const { status: finalStatus, exitCode, runtimeMs, failureType, failureClass, nextAction } = await attachToJob(config, dep.deployment_id, {
     chalk,
     maxRuntimeMs,
     maxCost,
@@ -853,11 +926,13 @@ export async function runCommand(config, args, chalk) {
     } else {
       console.error(formatCliError('JOB_FAILED', { exitCode, deploymentId: dep.deployment_id }, chalk));
     }
+    _printFailureClass(chalk, { failure_class: failureClass, next_action: nextAction });
     console.log(chalk.dim(_stage(STAGE_TOTAL, STAGE_TOTAL, 'Failed')));
     console.log(chalk.dim(`  Runtime: ${fmtRuntime(runtimeMs)}  •  Estimated cost: ~$${finalCost.toFixed(2)}`));
     // The container already reached a terminal state on the provider side by the
     // time we observe it here, so billing is already stopped — no extra teardown call needed.
     _printFinalInfo(chalk, { exitCode, teardownOk: true, jobId: dep.deployment_id, rcptId });
+    if (flags.resumeCmd) console.log(`  ${chalk.bold('Resume:')}    ${flags.resumeCmd}`);
     process.exitCode = exitCode ?? 1;
     return;
   }

package/src/commands/serve.js

@@ -4,7 +4,7 @@ import { addDeployment, addReceipt, updateReceipt, generateReceiptId } from '../
 import { normalizeTier, callWithFallback } from '../fallback.js';
 import { formatCliError } from '../errors.js';
 import { TEMPLATE_MAP, buildTemplateFlags, parseTemplateOverrides, BLESSED_VLLM_MODELS, isLikelyGatedModel } from '../catalog.js';
-import { stage as _stage, stageDone as _stageDone, writeBlock as _writeBlock, clearBlock as _clearBlock, renderLiveBlock as _renderLiveBlock } from '../progress.js';
+import { stage as _stage, stageDone as _stageDone, writeBlock as _writeBlock, clearBlock as _clearBlock, renderLiveBlock as _renderLiveBlock, printFailureClass as _printFailureClass } from '../progress.js';
 
 const LLAMA_CPP_IMAGE = 'michaelmanleyx/llama-cpp:server-cuda';
 
@@ -32,6 +32,7 @@ export function parseServeArgs(args) {
     if (args[i] === '--name')       { flags.name = args[++i]; i++; continue; }
     if (args[i] === '--no-wait')               { flags.noWait = true; i++; continue; }
     if (args[i] === '--max-cost')              { flags.maxCost = parseFloat(args[++i]); i++; continue; }
+    if (args[i] === '--idle-timeout')          { flags.idleTimeout = parseInt(args[++i], 10); i++; continue; }
     if (args[i] === '--health-path')           { flags.healthPath = args[++i]; i++; continue; }
     if (args[i] === '--check-nodes')           { flags.checkNodes = args[++i]; i++; continue; }
     // All three aliases map to noMarketplaceFallback
@@ -150,7 +151,12 @@ async function waitForEndpoint(deploymentId, config, timeoutMs = VLLM_SERVE_WAIT
       });
       if (['failed', 'terminated', 'error', 'stopped'].includes(dep.status)) {
         if (blockLines > 0) { _clearBlock(blockLines); blockLines = 0; }
-        return { ready: false, timedOut: false, depFailed: true, failReason: dep.fix_hint || dep.error || dep.status };
+        return {
+          ready: false, timedOut: false, depFailed: true,
+          failReason: dep.fix_hint || dep.error || dep.status,
+          failureClass: dep.failure_class ?? null,
+          nextAction: dep.next_action ?? null,
+        };
       }
       if (dep.endpoint_ready) {
         if (blockLines > 0) { _clearBlock(blockLines); blockLines = 0; }
@@ -203,11 +209,106 @@ async function validateComfyNodes(baseUrl, nodeList, chalk) {
 // Known badgr serve flags — used to detect broken shell line continuation.
 const _KNOWN_SERVE_FLAGS = new Set([
   '--gpu', '--image', '--task', '--count', '--region', '--tier', '--max-price',
-  '--name', '--no-wait', '--max-cost', '--health-path', '--check-nodes',
+  '--name', '--no-wait', '--max-cost', '--idle-timeout', '--health-path', '--check-nodes',
   '--no-fallback', '--strict-capacity', '--no-expanded-search', '--env',
   '--persistent', '--yes', '-y', '--runtime', '--hf-repo', '--hf-file',
 ]);
 
+function extractFlag(args, flagName) {
+  const rest = [];
+  let value = null;
+  for (let i = 0; i < args.length; i++) {
+    if (args[i] === flagName) { value = args[++i]; continue; }
+    rest.push(args[i]);
+  }
+  return { value, rest };
+}
+
+// Returns { id, endpointUrl } for a running vLLM endpoint matching modelId, or
+// null. The id is needed so a caller that auto-launches a new one can print a
+// precise `badgr down <id>` hint if something later fails.
+async function findRunningVllmEndpoint(config, modelId) {
+  try {
+    const existing = await listDeployments(config);
+    const ACTIVE = new Set(['running', 'provisioning', 'starting', 'queued']);
+    const match = (existing.deployments ?? []).find(d =>
+      ACTIVE.has(d.status) && d.workload_type === 'endpoint' && d.model === modelId
+    );
+    if (!match) return null;
+    const endpointUrl = match.endpoint_url || match.openai_base_url || null;
+    return endpointUrl ? { id: match.deployment_id, endpointUrl } : null;
+  } catch {
+    return null;
+  }
+}
+
+/**
+ * `badgr serve openwebui [--model <alias-or-id>] [--connect <url>] [flags]`
+ *
+ * Open WebUI serves a chat UI that connects to a model endpoint — usually
+ * vLLM. It is a separate served app from the model itself: this reuses a
+ * running vLLM endpoint for the model if one exists, or launches one via
+ * `serveCommand` first, then launches Open WebUI wired to it via
+ * OPENAI_API_BASE_URL/OPENAI_API_KEY. `--connect` skips vLLM discovery/launch
+ * entirely and points Open WebUI at any existing OpenAI-compatible endpoint —
+ * no model is ever touched in that mode.
+ */
+async function serveOpenWebUICommand(config, args, chalk) {
+  requireApiKey(config);
+  const { value: connect, rest: afterConnect } = extractFlag(args, '--connect');
+  const { value: model, rest } = extractFlag(afterConnect, '--model');
+
+  let endpointUrl = connect || null;
+  let launchedVllmDepId = null; // set only if we launched a new vLLM (for cleanup hints)
+
+  if (!endpointUrl) {
+    const alias = model || 'qwen-7b';
+    const vllmAlias = BLESSED_VLLM_MODELS[alias];
+    const modelId = vllmAlias ? vllmAlias.model_id : alias;
+
+    const existing = await findRunningVllmEndpoint(config, modelId);
+    if (existing) {
+      endpointUrl = existing.endpointUrl;
+      console.log(chalk.dim(`  Reusing running vLLM endpoint for ${modelId}: ${endpointUrl}\n`));
+    } else {
+      // Auto-launching vLLM reuses rest's --max-cost/--persistent, which also
+      // apply to Open WebUI below — that's two separately billed deployments,
+      // not a shared budget, so say so instead of a silent double spend.
+      const { value: webuiMaxCost } = extractFlag(rest, '--max-cost');
+      console.log(chalk.bold(`\n  Open WebUI needs a model endpoint behind it — no running vLLM endpoint for ${modelId} found.\n`));
+      if (webuiMaxCost) {
+        console.log(chalk.yellow(`  Launching vLLM first with its own --max-cost $${webuiMaxCost} cap (separate from Open WebUI's).`));
+        console.log(chalk.yellow(`  Total possible spend across both deployments: ~$${(parseFloat(webuiMaxCost) * 2).toFixed(2)}.\n`));
+      }
+      await serveCommand(config, [alias, ...rest], chalk);
+      const launched = await findRunningVllmEndpoint(config, modelId);
+      if (!launched) {
+        console.error(chalk.red(`\n  ✗ vLLM endpoint for ${modelId} did not come up — see output above.\n`));
+        console.error(chalk.dim(`  Once it's ready, run: badgr serve openwebui --connect <endpoint-url>\n`));
+        process.exitCode = 1;
+        return;
+      }
+      endpointUrl = launched.endpointUrl;
+      launchedVllmDepId = launched.id;
+    }
+  }
+
+  const overrides = parseTemplateOverrides(rest);
+  // User-supplied --env intentionally wins over the auto-wired connection env.
+  overrides.env = {
+    OPENAI_API_BASE_URL: endpointUrl,
+    OPENAI_API_KEY:      config.apiKey || 'sk-local',
+    ...overrides.env,
+  };
+  const flags = buildTemplateFlags(TEMPLATE_MAP.openwebui, overrides);
+  await serveCommand(config, flags, chalk);
+
+  if (process.exitCode && launchedVllmDepId) {
+    console.error(chalk.yellow(`\n  Open WebUI failed to start. The vLLM endpoint it was connecting to is still running and billing separately.`));
+    console.error(chalk.dim(`  Stop it with: badgr down ${launchedVllmDepId}\n`));
+  }
+}
+
 export async function serveCommand(config, args, chalk) {
   // `badgr serve template <name> [flags]` — expand template defaults then re-dispatch
   if (args[0] === 'template') {
@@ -230,6 +331,12 @@ export async function serveCommand(config, args, chalk) {
     return serveCommand(config, expandedArgs, chalk);
   }
 
+  // `badgr serve openwebui|open-webui [flags]` — a chat UI that connects to a
+  // model endpoint (usually vLLM). See serveOpenWebUICommand.
+  if (args[0] && ['openwebui', 'open-webui'].includes(args[0].toLowerCase())) {
+    return serveOpenWebUICommand(config, args.slice(1), chalk);
+  }
+
   if (args.includes('--list-aliases')) {
     console.log(chalk.bold('\nTested model routes:\n'));
     for (const alias of Object.keys(BLESSED_VLLM_MODELS)) {
@@ -437,6 +544,7 @@ export async function serveCommand(config, args, chalk) {
       tier: tierOverride || effectiveTier,
       ...(Object.keys(effectiveEnv).length > 0 ? { env: effectiveEnv } : {}),
       ...(flags.maxCost ? { max_cost_usd: flags.maxCost } : {}),
+      ...(flags.idleTimeout ? { idle_timeout_minutes: flags.idleTimeout } : {}),
       ...(flags.healthPath ? { health_path: flags.healthPath } : {}),
     };
   }
@@ -536,11 +644,12 @@ export async function serveCommand(config, args, chalk) {
 
   // Shared reporting for "deployment failed to start" — hit from both the
   // pre-poll status check and the waitForEndpoint poll loop below.
-  function reportDeployFailure(failReason) {
+  function reportDeployFailure(failReason, failureInfo = {}) {
     console.error(formatCliError('HEALTH_CHECK_DEPLOY_FAILED', {
       deploymentId: dep.deployment_id,
       failReason,
     }, chalk));
+    _printFailureClass(chalk, { failure_class: failureInfo.failureClass, next_action: failureInfo.nextAction });
     if (gatedHintNeeded) {
       const rerun = args.join(' ');
       const retryCmd = rerun.includes('HF_TOKEN=') ? rerun : `${rerun} --env HF_TOKEN=$HF_TOKEN`;
@@ -578,7 +687,9 @@ export async function serveCommand(config, args, chalk) {
         timeoutMs: 10_000,
       });
       if (['failed', 'terminated', 'error'].includes(latest.status)) {
-        reportDeployFailure(latest.error || latest.status);
+        reportDeployFailure(latest.error || latest.status, {
+          failureClass: latest.failure_class ?? null, nextAction: latest.next_action ?? null,
+        });
         return;
       }
     } catch {
@@ -592,7 +703,9 @@ export async function serveCommand(config, args, chalk) {
     );
 
     if (healthResult.depFailed) {
-      reportDeployFailure(healthResult.failReason);
+      reportDeployFailure(healthResult.failReason, {
+        failureClass: healthResult.failureClass, nextAction: healthResult.nextAction,
+      });
       return;
     }
 
@@ -635,22 +748,34 @@ export async function serveCommand(config, args, chalk) {
   else if (dep.model || effectiveModel) console.log(`  ${chalk.bold('Model:')}        ${dep.model || effectiveModel}`);
   if (customImage)        console.log(`  ${chalk.bold('Image:')}        ${customImage}`);
   if (flags.maxCost) console.log(`  ${chalk.bold('Max cost:')}     $${flags.maxCost.toFixed(2)}`);
+  if (flags.idleTimeout) console.log(`  ${chalk.bold('Idle timeout:')} ${flags.idleTimeout}m (auto-stops if idle — see Heartbeat below)`);
   console.log(`  ${chalk.bold('Logs:')}         badgr logs ${dep.deployment_id}`);
+  console.log(`  ${chalk.bold('Restart:')}      badgr restart ${dep.deployment_id}`);
   console.log(`  ${chalk.bold('Stop billing:')} ${chalk.cyan(`badgr down ${dep.deployment_id}`)}`);
   console.log(`  ${chalk.bold('Receipt:')}      badgr receipts ${rcptId}`);
   console.log();
   console.log(chalk.dim('  Billing continues until you run: ') + chalk.cyan(`badgr down ${dep.deployment_id}`));
 
   if (endpointReady && !customImage) {
-    const keySnip = config.apiKey?.slice(0, 4) || 'sk-...';
+    // dep.endpoint_api_key is a per-endpoint key generated for this deployment
+    // (vLLM model serves only) — shown exactly once, here. Falls back to the
+    // account-wide key (truncated) for serves that don't get one yet
+    // (custom images, managed transcribe/image tasks).
+    const hasEndpointKey = Boolean(dep.endpoint_api_key);
+    const authKey = hasEndpointKey ? dep.endpoint_api_key : `${config.apiKey?.slice(0, 4) || 'sk-...'}...`;
     const sdkModel = isLlamaCpp ? 'default' : (dep.model || effectiveModel);
+
+    if (hasEndpointKey) {
+      console.log(`  ${chalk.bold('API key:')}      ${chalk.yellow(dep.endpoint_api_key)}`);
+      console.log(chalk.dim('    Shown once — copy it now. This key is scoped to this endpoint only.'));
+    }
     console.log(`  ${chalk.bold('Test with curl:')}`);
     console.log(chalk.dim(`    curl ${endpointUrl}/chat/completions \\`));
-    console.log(chalk.dim(`      -H "Authorization: Bearer ${keySnip}..." -H "Content-Type: application/json" \\`));
+    console.log(chalk.dim(`      -H "Authorization: Bearer ${authKey}" -H "Content-Type: application/json" \\`));
     console.log(chalk.dim(`      -d '{"model":"${sdkModel}","messages":[{"role":"user","content":"Hello"}]}'`));
     console.log(`  ${chalk.bold('Use with OpenAI SDK:')}`);
     console.log(chalk.dim(`    from openai import OpenAI`));
-    console.log(chalk.dim(`    client = OpenAI(base_url="${endpointUrl}", api_key="${keySnip}...")`));
+    console.log(chalk.dim(`    client = OpenAI(base_url="${endpointUrl}", api_key="${authKey}")`));
     if (flags.task === 'transcribe') {
       console.log(chalk.dim(`    with open("audio.mp3", "rb") as f:`));
       console.log(chalk.dim(`        t = client.audio.transcriptions.create(model="${sdkModel}", file=f, response_format="text")`));
@@ -663,5 +788,13 @@ export async function serveCommand(config, args, chalk) {
       console.log(chalk.dim(`    resp = client.chat.completions.create(model="${sdkModel}", messages=[{"role": "user", "content": "Hello"}])`));
     }
     console.log();
+    if (flags.idleTimeout) {
+      console.log(`  ${chalk.bold('Heartbeat (required for --idle-timeout):')}`);
+      console.log(chalk.dim(`    badgr heartbeat ${dep.deployment_id}`));
+      console.log(chalk.dim('    Badgr does not proxy your inference traffic, so call this on each real'));
+      console.log(chalk.dim(`    request (or wire it into your client) — otherwise this endpoint auto-stops`));
+      console.log(chalk.dim(`    after ${flags.idleTimeout}m even if it's still reachable.`));
+      console.log();
+    }
   }
 }

package/src/commands/train.js

@@ -10,6 +10,7 @@ import { requireApiKey } from '../config.js';
 import { addReceipt, updateReceipt, generateReceiptId } from '../store.js';
 import { normalizeTier, callWithFallback } from '../fallback.js';
 import { monitorBatchJob, fmtRuntime } from '../batch.js';
+import { pollJobUntilTerminal, renderJobClosingBlock } from '../progress.js';
 
 const MAX_CONFIG_B = 512 * 1024; // 512 KB config limit
 
@@ -239,35 +240,29 @@ export async function trainLoraCommand(config, args, chalk) {
   console.log(`  ${chalk.bold('Receipt:')} ${chalk.dim(rcptId)}`);
   console.log(chalk.dim('\n  Polling for completion — Ctrl+C to detach (GPU keeps running)\n'));
 
-  // Poll until complete
-  const startMs = Date.now();
   const maxMs = maxRuntime * 60 * 1000;
-  while (Date.now() - startMs < maxMs) {
-    await new Promise(r => setTimeout(r, 15_000));
-    let detail;
-    try {
-      detail = await callApi(`/jobs/${job.job_id}`, {
-        apiKey: config.apiKey,
-        baseUrl: config.baseUrl,
-      });
-    } catch { continue; }
-    process.stdout.write(`\r  Status: ${detail.status}  elapsed: ${Math.floor((Date.now() - startMs) / 1000)}s  `);
-    if (detail.status === 'completed') {
-      const out = detail.output || {};
-      console.log(chalk.green('\n\n  ✓ Training complete\n'));
-      if (out.adapter_url) console.log(`  ${chalk.bold('Adapter:')} ${out.adapter_url}`);
-      if (out.checkpoint_url) console.log(`  ${chalk.bold('Checkpoint:')} ${out.checkpoint_url}`);
-      console.log(`  ${chalk.bold('Receipt:')} ${chalk.dim(rcptId)}\n`);
-      return;
-    }
-    if (detail.status === 'failed') {
-      console.error(chalk.red(`\n\n  ✗ Training failed: ${detail.error_code || ''} — ${detail.error_message || ''}\n`));
-      process.exitCode = 1;
-      return;
-    }
+  const { outcome, detail } = await pollJobUntilTerminal(callApi, config, job.job_id, { chalk, maxMs });
+
+  if (outcome === 'polling_failed') {
+    console.error(chalk.red(`\n\n  ✗ Lost contact with Badgr — could not confirm job status.\n`));
+    console.error(chalk.dim(`    badgr status\n    badgr logs ${job.job_id}\n`));
+    process.exitCode = 1;
+    return;
+  }
+  if (outcome === 'timed_out') {
+    console.error(chalk.yellow('\n\n  Training still running — detached. Check status:\n'));
+    console.error(chalk.dim(`    badgr status\n`));
+    return;
+  }
+
+  if (outcome === 'completed') {
+    const out = detail.output || {};
+    if (out.adapter_url) console.log(`\n  ${chalk.bold('Adapter:')} ${out.adapter_url}`);
+    if (out.checkpoint_url) console.log(`  ${chalk.bold('Checkpoint:')} ${out.checkpoint_url}`);
+    if (out.warning) console.log(chalk.yellow(`  ${out.warning}`));
   }
-  console.error(chalk.yellow('\n\n  Training still running — detached. Check status:\n'));
-  console.error(chalk.dim(`    badgr status\n`));
+  console.log(renderJobClosingBlock(chalk, detail, rcptId));
+  if (outcome === 'failed') process.exitCode = 1;
 }
 
 export async function trainCommand(config, args, chalk) {