badgr-cli 1.0.45 → 1.0.46

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "badgr-cli",
-  "version": "1.0.45",
+  "version": "1.0.46",
   "description": "Badgr — run or serve GPU workloads from one command",
   "type": "module",
   "bin": {

package/src/api.js

@@ -168,6 +168,24 @@ export async function terminateDeployment(config, deploymentId) {
   throw lastErr;
 }
 
+export function restartDeployment(config, deploymentId) {
+  return callApi(`/deployments/${deploymentId}/restart`, {
+    method: 'POST',
+    apiKey: config.apiKey,
+    baseUrl: config.baseUrl,
+    timeoutMs: 30_000,
+  });
+}
+
+export function heartbeatDeployment(config, deploymentId) {
+  return callApi(`/deployments/${deploymentId}/heartbeat`, {
+    method: 'POST',
+    apiKey: config.apiKey,
+    baseUrl: config.baseUrl,
+    timeoutMs: 10_000,
+  });
+}
+
 export function getDeploymentLogs(config, deploymentId) {
   return callApi(`/deployments/${deploymentId}/logs`, {
     apiKey: config.apiKey,

package/src/badgr.js

@@ -21,6 +21,8 @@ import { templateCommand }   from './commands/template.js';
 import { workloadCommand }   from './commands/workload.js';
 import { workspaceCommand }  from './commands/workspace.js';
 import { detectCommand }     from './commands/detect.js';
+import { restartCommand }    from './commands/restart.js';
+import { heartbeatCommand }  from './commands/heartbeat.js';
 
 const HELP = `
 ${chalk.bold('badgr')} — run or serve GPU workloads from one command
@@ -34,6 +36,8 @@ ${chalk.bold('COMMANDS')}
   ${chalk.cyan('badgr status')}              Show what's running and what's billing
   ${chalk.cyan('badgr logs <id>')}           Stream logs for a running job or endpoint
   ${chalk.cyan('badgr down <id>')}           Stop a deployment and end billing
+  ${chalk.cyan('badgr restart <id>')}        Relaunch an endpoint with the same config and API key
+  ${chalk.cyan('badgr heartbeat <id>')}      Reset an endpoint's idle-timeout clock
   ${chalk.cyan('badgr receipts')}            Show cost history
   ${chalk.cyan('badgr test')}               Run an end-to-end test (provision → run → teardown)
   ${chalk.cyan('badgr capacity')}            Check what GPU capacity is available right now
@@ -158,6 +162,8 @@ async function main() {
     case 'status':   return statusCommand(config, rest, chalk);
     case 'logs':     return logsCommand(config, rest, chalk);
     case 'down':     return downCommand(config, rest, chalk);
+    case 'restart':   return restartCommand(config, rest, chalk);
+    case 'heartbeat': return heartbeatCommand(config, rest, chalk);
     case 'receipts': return receiptsCommand(config, rest, chalk);
     case 'models':   return modelsCommand(config, chalk);
     case 'capacity': return capacityCommand(config, rest, chalk);

package/src/commands/heartbeat.js

@@ -0,0 +1,38 @@
+import { requireApiKey } from '../config.js';
+import { findDeployment } from '../store.js';
+import { heartbeatDeployment } from '../api.js';
+
+/**
+ * badgr heartbeat <deployment-id|name>
+ *
+ * Resets an endpoint's idle-timeout clock. Badgr doesn't proxy inference
+ * traffic, so if you set --idle-timeout on `badgr serve`, call this on
+ * each real request (or wire it into your own client) — otherwise the
+ * endpoint is torn down once idle_timeout_minutes elapses, even if it's
+ * still reachable.
+ */
+export async function heartbeatCommand(config, args, chalk) {
+  const idOrName = args.find(a => !a.startsWith('--'));
+
+  if (!idOrName) {
+    console.error(chalk.red('Usage: badgr heartbeat <deployment-id|name>'));
+    console.error(chalk.dim('       Resets the idle-timeout clock for an endpoint launched with --idle-timeout.'));
+    return;
+  }
+
+  requireApiKey(config);
+
+  const localDep     = findDeployment(idOrName);
+  const deploymentId = localDep?.id ?? idOrName;
+
+  try {
+    const result = await heartbeatDeployment(config, deploymentId);
+    console.log(chalk.green(`  ✓ Heartbeat recorded for ${deploymentId}`));
+    if (result.last_activity_at) {
+      console.log(chalk.dim(`    last_activity_at: ${new Date(result.last_activity_at * 1000).toISOString()}`));
+    }
+  } catch (err) {
+    console.error(chalk.red(`  ✗ Could not send heartbeat: ${err.message}`));
+    process.exitCode = 1;
+  }
+}

package/src/commands/restart.js

@@ -0,0 +1,74 @@
+import { requireApiKey } from '../config.js';
+import { findDeployment, removeDeployment, addDeployment, addReceipt, generateReceiptId } from '../store.js';
+import { restartDeployment } from '../api.js';
+
+/**
+ * badgr restart <deployment-id|name>
+ *
+ * Tears down the current pod and relaunches an endpoint with the same
+ * config (GPU, model, price/cost/runtime caps, endpoint API key). Returns
+ * a *new* deployment_id and endpoint_url — the old pod's IP is gone.
+ */
+export async function restartCommand(config, args, chalk) {
+  const idOrName = args.find(a => !a.startsWith('--'));
+
+  if (!idOrName) {
+    console.error(chalk.red('Usage: badgr restart <deployment-id|name>'));
+    return;
+  }
+
+  requireApiKey(config);
+
+  const localDep     = findDeployment(idOrName);
+  const deploymentId = localDep?.id ?? idOrName;
+
+  process.stdout.write(chalk.dim(`  Restarting ${deploymentId}...`));
+
+  let dep;
+  try {
+    dep = await restartDeployment(config, deploymentId);
+  } catch (err) {
+    process.stdout.write('\n');
+    console.error(chalk.red(`\n  ✗ Could not restart deployment: ${err.message}\n`));
+    process.exitCode = 1;
+    return;
+  }
+
+  process.stdout.write('\n');
+
+  removeDeployment(idOrName);
+  addDeployment({
+    id:          dep.deployment_id,
+    name:        dep.name,
+    type:        dep.workload_type,
+    model:       dep.model,
+    gpu:         dep.gpu_type,
+    count:       dep.gpu_count,
+    status:      dep.status,
+    endpointUrl: dep.endpoint_url || dep.openai_base_url,
+    receiptId:   dep.receipt_id,
+    createdAt:   new Date().toISOString(),
+    costPerHour: dep.cost_per_hour || 0,
+    providerRoute: dep.provider ?? null,
+    tier:          dep.tier ?? null,
+  });
+
+  const rcptId = dep.receipt_id || generateReceiptId();
+  addReceipt({
+    receiptId:    rcptId,
+    action:       'badgr restart',
+    deploymentId: dep.deployment_id,
+    gpu:          dep.gpu_type,
+    status:       dep.status,
+    createdAt:    new Date().toISOString(),
+  });
+
+  const endpointUrl = dep.endpoint_url || dep.openai_base_url;
+  console.log(chalk.green('\n  ✓ Restarted\n'));
+  console.log(`  ${chalk.bold('New deployment:')} ${chalk.cyan(dep.deployment_id)}`);
+  if (endpointUrl) console.log(`  ${chalk.bold('New base URL:')}   ${chalk.cyan(endpointUrl)}`);
+  console.log(chalk.dim('  Old endpoint URL is gone — update any client pointed at it.'));
+  console.log(chalk.dim('  Your endpoint API key is unchanged — no need to re-issue it.'));
+  console.log(`  ${chalk.bold('Status:')}         badgr status ${dep.deployment_id}`);
+  console.log();
+}

package/src/commands/serve.js

@@ -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
@@ -208,7 +209,7 @@ 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',
 ]);
@@ -543,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 } : {}),
     };
   }
@@ -746,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")`));
@@ -774,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/tests/heartbeat.test.js

@@ -0,0 +1,70 @@
+/**
+ * badgr heartbeat — resets an endpoint's idle-timeout clock
+ */
+import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
+import { heartbeatCommand } from '../src/commands/heartbeat.js';
+
+vi.mock('../src/api.js', () => ({
+  heartbeatDeployment: vi.fn(),
+}));
+
+vi.mock('../src/store.js', () => ({
+  findDeployment: vi.fn(() => null),
+}));
+
+vi.mock('../src/config.js', () => ({
+  requireApiKey: vi.fn(),
+}));
+
+import * as api   from '../src/api.js';
+import * as store from '../src/store.js';
+
+const config = { apiKey: 'sk-test', baseUrl: 'https://api.test/v1' };
+
+const chalk = {
+  bold:   s => s,
+  dim:    s => s,
+  red:    s => s,
+  yellow: s => s,
+  green:  s => s,
+  cyan:   s => s,
+};
+
+beforeEach(() => {
+  process.exitCode = undefined;
+  vi.spyOn(console, 'log').mockImplementation(() => {});
+  vi.spyOn(console, 'error').mockImplementation(() => {});
+  vi.resetAllMocks();
+  store.findDeployment.mockReturnValue(null);
+});
+
+afterEach(() => {
+  vi.restoreAllMocks();
+  process.exitCode = undefined;
+});
+
+it('requires a deployment id', async () => {
+  await heartbeatCommand(config, [], chalk);
+  expect(api.heartbeatDeployment).not.toHaveBeenCalled();
+});
+
+it('calls heartbeatDeployment and prints confirmation', async () => {
+  api.heartbeatDeployment.mockResolvedValue({ deployment_id: 'dep-abc123', last_activity_at: 1700000000 });
+
+  const lines = [];
+  console.log.mockImplementation((...args) => lines.push(args.join(' ')));
+
+  await heartbeatCommand(config, ['dep-abc123'], chalk);
+
+  expect(api.heartbeatDeployment).toHaveBeenCalledWith(config, 'dep-abc123');
+  expect(lines.some(l => l.includes('Heartbeat recorded for dep-abc123'))).toBe(true);
+  expect(process.exitCode).toBeFalsy();
+});
+
+it('reports an error and sets exitCode on failure', async () => {
+  api.heartbeatDeployment.mockRejectedValue(new Error('deployment not found'));
+
+  await heartbeatCommand(config, ['dep-missing'], chalk);
+
+  expect(process.exitCode).toBe(1);
+});

package/tests/restart.test.js

@@ -0,0 +1,88 @@
+/**
+ * badgr restart — relaunches an endpoint with the same config and API key
+ */
+import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
+import { restartCommand } from '../src/commands/restart.js';
+
+vi.mock('../src/api.js', () => ({
+  restartDeployment: vi.fn(),
+}));
+
+vi.mock('../src/store.js', () => ({
+  findDeployment:    vi.fn(() => null),
+  removeDeployment:  vi.fn(),
+  addDeployment:     vi.fn(),
+  addReceipt:        vi.fn(),
+  generateReceiptId: vi.fn(() => 'rcpt-restart-001'),
+}));
+
+vi.mock('../src/config.js', () => ({
+  requireApiKey: vi.fn(),
+}));
+
+import * as api   from '../src/api.js';
+import * as store from '../src/store.js';
+
+const config = { apiKey: 'sk-test', baseUrl: 'https://api.test/v1' };
+
+const chalk = {
+  bold:   s => s,
+  dim:    s => s,
+  red:    s => s,
+  yellow: s => s,
+  green:  s => s,
+  cyan:   s => s,
+};
+
+beforeEach(() => {
+  process.exitCode = undefined;
+  vi.spyOn(console, 'log').mockImplementation(() => {});
+  vi.spyOn(console, 'error').mockImplementation(() => {});
+  vi.spyOn(process.stdout, 'write').mockImplementation(() => true);
+  vi.resetAllMocks();
+  store.findDeployment.mockReturnValue(null);
+  store.generateReceiptId.mockReturnValue('rcpt-restart-001');
+});
+
+afterEach(() => {
+  vi.restoreAllMocks();
+  process.exitCode = undefined;
+});
+
+it('requires a deployment id', async () => {
+  await restartCommand(config, [], chalk);
+  expect(api.restartDeployment).not.toHaveBeenCalled();
+});
+
+it('calls restartDeployment and prints the new deployment id and URL', async () => {
+  api.restartDeployment.mockResolvedValue({
+    deployment_id: 'dep-new-002',
+    workload_type: 'endpoint',
+    model: 'Qwen/Qwen2.5-7B-Instruct',
+    gpu_type: 'L40S',
+    gpu_count: 1,
+    status: 'provisioning',
+    endpoint_url: 'https://dep-new-002.aibadgr.com/v1',
+    cost_per_hour: 1.2,
+  });
+
+  const lines = [];
+  console.log.mockImplementation((...args) => lines.push(args.join(' ')));
+
+  await restartCommand(config, ['dep-old-001'], chalk);
+
+  expect(api.restartDeployment).toHaveBeenCalledWith(config, 'dep-old-001');
+  expect(store.removeDeployment).toHaveBeenCalledWith('dep-old-001');
+  expect(store.addDeployment).toHaveBeenCalledWith(expect.objectContaining({ id: 'dep-new-002' }));
+  expect(lines.some(l => l.includes('dep-new-002'))).toBe(true);
+  expect(lines.some(l => l.includes('https://dep-new-002.aibadgr.com/v1'))).toBe(true);
+  expect(process.exitCode).toBeFalsy();
+});
+
+it('reports an error and sets exitCode on failure', async () => {
+  api.restartDeployment.mockRejectedValue(new Error('deployment not found'));
+
+  await restartCommand(config, ['dep-missing'], chalk);
+
+  expect(process.exitCode).toBe(1);
+});

package/tests/serve-lifecycle.test.js

@@ -425,6 +425,78 @@ describe('custom image health check', () => {
     expect(process.exitCode).toBeFalsy();
   });
 
+  it('sends --idle-timeout to the backend as idle_timeout_minutes', async () => {
+    api.callApi
+      .mockResolvedValueOnce(makeServeDep())
+      .mockResolvedValueOnce({ status: 'running' })
+      .mockResolvedValueOnce(readyStatus());
+
+    const p = serveCommand(
+      config,
+      ['meta-llama/Llama-3.1-8B-Instruct', '--idle-timeout', '30', '--max-cost', '5'],
+      chalk,
+    );
+    await vi.advanceTimersByTimeAsync(5000);
+    await p;
+
+    const body = api.callApi.mock.calls[0][1].body;
+    expect(body.idle_timeout_minutes).toBe(30);
+    expect(process.exitCode).toBeFalsy();
+  });
+
+  it('omits idle_timeout_minutes when --idle-timeout is not passed', async () => {
+    api.callApi
+      .mockResolvedValueOnce(makeServeDep())
+      .mockResolvedValueOnce({ status: 'running' })
+      .mockResolvedValueOnce(readyStatus());
+
+    const p = serveCommand(config, ['meta-llama/Llama-3.1-8B-Instruct', '--max-cost', '5'], chalk);
+    await vi.advanceTimersByTimeAsync(5000);
+    await p;
+
+    const body = api.callApi.mock.calls[0][1].body;
+    expect(body.idle_timeout_minutes).toBeUndefined();
+  });
+
+  it('prints the per-endpoint API key exactly once when the backend returns one', async () => {
+    api.callApi
+      .mockResolvedValueOnce(makeServeDep({ endpoint_api_key: 'bge_test_key_xyz' }))
+      .mockResolvedValueOnce({ status: 'running' })
+      .mockResolvedValueOnce(readyStatus());
+
+    const lines = [];
+    console.log.mockImplementation((...args) => lines.push(args.join(' ')));
+
+    const p = serveCommand(config, ['meta-llama/Llama-3.1-8B-Instruct', '--max-cost', '5'], chalk);
+    await vi.advanceTimersByTimeAsync(5000);
+    await p;
+
+    const keyLine = lines.find(l => l.includes('API key:'));
+    expect(keyLine).toContain('bge_test_key_xyz');
+    const sdkLine = lines.find(l => l.includes('api_key='));
+    expect(sdkLine).toContain('bge_test_key_xyz');
+  });
+
+  it('prints a heartbeat reminder when --idle-timeout is set', async () => {
+    api.callApi
+      .mockResolvedValueOnce(makeServeDep({ endpoint_api_key: 'bge_test_key_xyz' }))
+      .mockResolvedValueOnce({ status: 'running' })
+      .mockResolvedValueOnce(readyStatus());
+
+    const lines = [];
+    console.log.mockImplementation((...args) => lines.push(args.join(' ')));
+
+    const p = serveCommand(
+      config,
+      ['meta-llama/Llama-3.1-8B-Instruct', '--idle-timeout', '15', '--max-cost', '5'],
+      chalk,
+    );
+    await vi.advanceTimersByTimeAsync(5000);
+    await p;
+
+    expect(lines.some(l => l.includes('badgr heartbeat dep-serve-001'))).toBe(true);
+  });
+
   it('--health-path also works for model-based serve (overrides /models default)', async () => {
     api.callApi
       .mockResolvedValueOnce(makeServeDep())                 // POST /serve