badgr-cli 1.0.45 → 1.0.47

package/tests/batch.test.js

@@ -0,0 +1,329 @@
+/**
+ * badgr batch — /run body shape, missing-artifact-path detection, and
+ * `badgr batch compare` regression logic against literal metrics.json
+ * fixtures (satisfies test scenario 5 without needing a live run).
+ */
+import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
+import { mkdtempSync, writeFileSync, mkdirSync, rmSync } from 'fs';
+import { tmpdir } from 'os';
+import { join } from 'path';
+
+vi.mock('../src/store.js', () => ({
+  addReceipt: vi.fn(),
+  updateReceipt: vi.fn(),
+  generateReceiptId: vi.fn(() => 'rcpt-batch-001'),
+  loadStore: vi.fn(() => ({ receipts: [], deployments: [] })),
+}));
+
+vi.mock('../src/api.js', () => ({
+  callApi: vi.fn(),
+  uploadBlob: vi.fn(),
+}));
+
+vi.mock('../src/fallback.js', async (importOriginal) => {
+  const actual = await importOriginal();
+  return { ...actual, callWithFallback: vi.fn() };
+});
+
+vi.mock('../src/batch.js', () => ({
+  monitorBatchJob: vi.fn(() => Promise.resolve({ status: 'succeeded', exitCode: 0, runtimeMs: 42000, reason: 'complete' })),
+  fmtRuntime: (ms) => `${Math.round(ms / 1000)}s`,
+}));
+
+vi.mock('archiver', () => ({
+  default: () => {
+    let out;
+    return {
+      on: () => {},
+      pipe: (o) => { out = o; },
+      directory: () => {},
+      file: () => {},
+      finalize: () => { out.end(); },
+    };
+  },
+}));
+
+import {
+  batchCommand,
+  displayFailureReason,
+  missingOutputPaths,
+  readMetricFile,
+} from '../src/commands/batch.js';
+import * as store from '../src/store.js';
+import * as api from '../src/api.js';
+import * as fallback from '../src/fallback.js';
+
+const chalk = {
+  bold: s => s, dim: s => s, red: s => s, yellow: s => s, green: s => s, cyan: s => s,
+};
+
+const config = { apiKey: 'sk-test', baseUrl: 'https://api.test/v1' };
+
+let dir;
+beforeEach(() => {
+  dir = mkdtempSync(join(tmpdir(), 'badgr-batch-cmd-test-'));
+  process.exitCode = undefined;
+  vi.spyOn(console, 'log').mockImplementation(() => {});
+  vi.spyOn(console, 'error').mockImplementation(() => {});
+  vi.spyOn(process.stdout, 'write').mockImplementation(() => true);
+  vi.clearAllMocks();
+  store.loadStore.mockReturnValue({ receipts: [], deployments: [] });
+});
+
+afterEach(() => {
+  rmSync(dir, { recursive: true, force: true });
+  vi.restoreAllMocks();
+});
+
+// ---------------------------------------------------------------------------
+// displayFailureReason
+// ---------------------------------------------------------------------------
+
+describe('displayFailureReason', () => {
+  it('maps max_runtime to runtime_limit_exceeded', () => {
+    expect(displayFailureReason('max_runtime')).toBe('runtime_limit_exceeded');
+  });
+  it('maps max_cost to cost_cap_exceeded', () => {
+    expect(displayFailureReason('max_cost')).toBe('cost_cap_exceeded');
+  });
+  it('passes through unknown reasons unchanged', () => {
+    expect(displayFailureReason('something_else')).toBe('something_else');
+  });
+  it('returns null for no reason', () => {
+    expect(displayFailureReason(null)).toBeNull();
+    expect(displayFailureReason(undefined)).toBeNull();
+  });
+});
+
+// ---------------------------------------------------------------------------
+// missingOutputPaths / readMetricFile
+// ---------------------------------------------------------------------------
+
+describe('missingOutputPaths', () => {
+  it('reports declared paths that do not exist under the extraction dir', () => {
+    mkdirSync(join(dir, 'outputs'), { recursive: true });
+    writeFileSync(join(dir, 'outputs', 'metrics.json'), '{}');
+    const missing = missingOutputPaths(dir, ['/outputs/metrics.json', '/outputs/videos', '/outputs/logs']);
+    expect(missing).toEqual(['/outputs/videos', '/outputs/logs']);
+  });
+
+  it('reports nothing missing when all declared paths exist', () => {
+    mkdirSync(join(dir, 'outputs', 'videos'), { recursive: true });
+    writeFileSync(join(dir, 'outputs', 'metrics.json'), '{}');
+    const missing = missingOutputPaths(dir, ['/outputs/metrics.json', '/outputs/videos']);
+    expect(missing).toEqual([]);
+  });
+});
+
+describe('readMetricFile', () => {
+  it('reads and parses a metrics JSON file', () => {
+    mkdirSync(join(dir, 'outputs'), { recursive: true });
+    writeFileSync(join(dir, 'outputs', 'metrics.json'), JSON.stringify({ pass_rate: 0.9 }));
+    expect(readMetricFile(dir, '/outputs/metrics.json')).toEqual({ pass_rate: 0.9 });
+  });
+
+  it('returns null when the file does not exist', () => {
+    expect(readMetricFile(dir, '/outputs/metrics.json')).toBeNull();
+  });
+
+  it('returns null on malformed JSON', () => {
+    mkdirSync(join(dir, 'outputs'), { recursive: true });
+    writeFileSync(join(dir, 'outputs', 'metrics.json'), 'not json');
+    expect(readMetricFile(dir, '/outputs/metrics.json')).toBeNull();
+  });
+});
+
+// ---------------------------------------------------------------------------
+// badgr batch run — /run body shape (image/command/env/caps)
+// ---------------------------------------------------------------------------
+
+describe('badgr batch run — /run request shape', () => {
+  it('builds the /run body from workload.yml with inputs/outputs env wiring', async () => {
+    writeFileSync(join(dir, 'policy.py'), 'def act(obs): return 0');
+    mkdirSync(join(dir, 'scenarios'));
+    writeFileSync(join(dir, 'scenarios', 's0.json'), '{}');
+    const yamlPath = join(dir, 'workload.yml');
+    writeFileSync(yamlPath, `
+name: toy-robot-eval
+image: badgr/physical-ai-eval-lite:latest
+command: [python, run_eval.py, --policy, /inputs/policy.py, --scenarios, /inputs/scenarios]
+inputs:
+  - ./policy.py:/inputs/policy.py
+  - ./scenarios:/inputs/scenarios
+outputs:
+  - /outputs/metrics.json
+  - /outputs/videos
+  - /outputs/logs
+max_cost: 20
+max_runtime_minutes: 60
+success_metric:
+  file: /outputs/metrics.json
+  key: pass_rate
+  higher_is_better: true
+`);
+
+    let capturedBody;
+    fallback.callWithFallback.mockImplementation((endpoint, callOpts, buildBody) => {
+      capturedBody = buildBody();
+      return Promise.resolve({
+        deployment_id: 'dep-abc123', receipt_id: 'rcpt-1', gpu_type: 'RTX_4090',
+        gpu_count: 1, cost_per_hour: 0.5, provider: 'runpod', tier: '1', status: 'running',
+      });
+    });
+
+    api.callApi.mockImplementation((path) => {
+      if (path.endsWith('/artifacts')) return Promise.reject(new Error('no artifacts yet'));
+      if (path.startsWith('/deployments/')) {
+        return Promise.resolve({
+          status: 'succeeded', failure_reason: null, teardown_ok: 'ok',
+          runtime_seconds: 42, accrued_cost_usd: 0.01, provider: 'runpod',
+        });
+      }
+      return Promise.resolve({});
+    });
+
+    api.uploadBlob.mockResolvedValue({ code_uri: 'https://cdn.example.com/inputs.tar.gz', upload_id: 'up_1' });
+
+    await batchCommand(config, ['run', yamlPath], chalk);
+
+    expect(capturedBody.image).toBe('badgr/physical-ai-eval-lite:latest');
+    expect(capturedBody.command).toEqual(['python', 'run_eval.py', '--policy', '/inputs/policy.py', '--scenarios', '/inputs/scenarios']);
+    expect(capturedBody.max_cost_usd).toBe(20);
+    expect(capturedBody.max_runtime_seconds).toBe(3600);
+    expect(capturedBody.custom_image).toBe(true);
+    expect(capturedBody.env.BADGR_INPUTS_ARCHIVE_URL).toBe('https://cdn.example.com/inputs.tar.gz');
+    expect(JSON.parse(capturedBody.env.BADGR_OUTPUT_PATHS_JSON)).toEqual([
+      '/outputs/metrics.json', '/outputs/videos', '/outputs/logs',
+    ]);
+  });
+
+  it('rejects a workload.yml missing required fields before touching the network', async () => {
+    const yamlPath = join(dir, 'bad.yml');
+    writeFileSync(yamlPath, 'image: busybox\n');
+    await batchCommand(config, ['run', yamlPath], chalk);
+    expect(process.exitCode).toBe(1);
+    expect(fallback.callWithFallback).not.toHaveBeenCalled();
+  });
+
+  it('builds a body with no BADGR_* env vars when inputs/outputs are omitted', async () => {
+    const yamlPath = join(dir, 'workload.yml');
+    writeFileSync(yamlPath, `
+name: simple-job
+image: busybox
+command: [echo, hi]
+max_cost: 1
+max_runtime_minutes: 5
+`);
+
+    let capturedBody;
+    fallback.callWithFallback.mockImplementation((endpoint, callOpts, buildBody) => {
+      capturedBody = buildBody();
+      return Promise.resolve({
+        deployment_id: 'dep-xyz', receipt_id: 'rcpt-2', gpu_type: 'RTX_4090',
+        gpu_count: 1, cost_per_hour: 0.1, provider: 'runpod', tier: '1', status: 'running',
+      });
+    });
+    api.callApi.mockImplementation((path) => {
+      if (path.startsWith('/deployments/')) {
+        return Promise.resolve({ status: 'succeeded', teardown_ok: 'ok' });
+      }
+      return Promise.resolve({});
+    });
+
+    await batchCommand(config, ['run', yamlPath], chalk);
+
+    expect(capturedBody.env.BADGR_INPUTS_ARCHIVE_URL).toBeUndefined();
+    expect(capturedBody.env.BADGR_OUTPUT_PATHS_JSON).toBeUndefined();
+  });
+});
+
+// ---------------------------------------------------------------------------
+// badgr batch artifacts — clear, actionable failure messages
+// ---------------------------------------------------------------------------
+
+describe('badgr batch artifacts — error clarity', () => {
+  it('explains a 404 (no artifact ever uploaded) with likely causes and a next step', async () => {
+    const err = new Error('GET /deployments/dep-nofiles/artifacts → HTTP 404: artifact_not_found');
+    err.httpStatus = 404;
+    api.callApi.mockImplementation((path) => {
+      if (path.endsWith('/artifacts')) return Promise.reject(err);
+      return Promise.resolve({});
+    });
+
+    await batchCommand(config, ['artifacts', 'dep-nofiles'], chalk);
+
+    expect(process.exitCode).toBe(1);
+    const logged = console.error.mock.calls.flat().join('\n');
+    expect(logged).toMatch(/No artifact was ever uploaded for run dep-nofiles/);
+    expect(logged).toMatch(/badgr batch logs dep-nofiles/);
+  });
+
+  it('explains a 410 (expired, 48h retention) plainly', async () => {
+    const err = new Error('GET /deployments/dep-old/artifacts → HTTP 410: artifact_expired');
+    err.httpStatus = 410;
+    api.callApi.mockImplementation((path) => {
+      if (path.endsWith('/artifacts')) return Promise.reject(err);
+      return Promise.resolve({});
+    });
+
+    await batchCommand(config, ['artifacts', 'dep-old'], chalk);
+
+    expect(process.exitCode).toBe(1);
+    const logged = console.error.mock.calls.flat().join('\n');
+    expect(logged).toMatch(/expired \(48h retention\)/);
+  });
+});
+
+// ---------------------------------------------------------------------------
+// badgr batch compare — direct metrics.json files (test scenario 5)
+// ---------------------------------------------------------------------------
+
+describe('badgr batch compare', () => {
+  it('reports a regression with delta when run B is worse (higher_is_better)', async () => {
+    const a = join(dir, 'a.json');
+    const b = join(dir, 'b.json');
+    writeFileSync(a, JSON.stringify({ pass_rate: 0.92 }));
+    writeFileSync(b, JSON.stringify({ pass_rate: 0.87 }));
+
+    await batchCommand(config, ['compare', a, b], chalk);
+
+    expect(process.exitCode).toBe(1);
+    const logged = console.log.mock.calls.flat().join('\n');
+    expect(logged).toMatch(/Regression/);
+    expect(logged).toContain('0.92');
+    expect(logged).toContain('0.87');
+  });
+
+  it('reports an improvement when run B is better', async () => {
+    const a = join(dir, 'a.json');
+    const b = join(dir, 'b.json');
+    writeFileSync(a, JSON.stringify({ pass_rate: 0.80 }));
+    writeFileSync(b, JSON.stringify({ pass_rate: 0.95 }));
+
+    await batchCommand(config, ['compare', a, b], chalk);
+
+    expect(process.exitCode).toBeUndefined();
+    const logged = console.log.mock.calls.flat().join('\n');
+    expect(logged).toMatch(/Improvement/);
+  });
+
+  it('respects --key and --higher-is-better overrides for raw files', async () => {
+    const a = join(dir, 'a.json');
+    const b = join(dir, 'b.json');
+    writeFileSync(a, JSON.stringify({ error_rate: 0.05 }));
+    writeFileSync(b, JSON.stringify({ error_rate: 0.10 }));
+
+    // error_rate: lower is better, so b (higher) is a regression.
+    await batchCommand(config, ['compare', a, b, '--key', 'error_rate', '--higher-is-better', 'false'], chalk);
+
+    expect(process.exitCode).toBe(1);
+    const logged = console.log.mock.calls.flat().join('\n');
+    expect(logged).toMatch(/Regression/);
+  });
+
+  it('errors out when a run_id has no known success_metric and no local YAML', async () => {
+    store.loadStore.mockReturnValueOnce({ receipts: [], deployments: [] });
+    await batchCommand(config, ['compare', 'dep-unknown-a', 'dep-unknown-b'], chalk);
+    expect(process.exitCode).toBe(1);
+  });
+});

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/rerun.test.js

@@ -0,0 +1,94 @@
+/**
+ * badgr rerun — replays a past job or endpoint with its exact original spec
+ */
+import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
+import { rerunCommand } from '../src/commands/rerun.js';
+
+vi.mock('../src/api.js', () => ({
+  rerunDeployment: vi.fn(),
+}));
+
+vi.mock('../src/store.js', () => ({
+  findDeployment:    vi.fn(() => null),
+  addDeployment:     vi.fn(),
+  addReceipt:        vi.fn(),
+  generateReceiptId: vi.fn(() => 'rcpt-rerun-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-rerun-001');
+});
+
+afterEach(() => {
+  vi.restoreAllMocks();
+  process.exitCode = undefined;
+});
+
+it('requires a deployment id', async () => {
+  await rerunCommand(config, [], chalk);
+  expect(api.rerunDeployment).not.toHaveBeenCalled();
+});
+
+it('calls rerunDeployment and prints the new deployment id, never removing the source', async () => {
+  api.rerunDeployment.mockResolvedValue({
+    deployment_id: 'dep-new-002',
+    workload_type: 'job',
+    gpu_type: 'A100',
+    gpu_count: 1,
+    status: 'provisioning',
+    rerun_of: 'dep-old-001',
+    cost_per_hour: 1.9,
+  });
+
+  const lines = [];
+  console.log.mockImplementation((...args) => lines.push(args.join(' ')));
+
+  await rerunCommand(config, ['dep-old-001'], chalk);
+
+  expect(api.rerunDeployment).toHaveBeenCalledWith(config, 'dep-old-001');
+  expect(store.addDeployment).toHaveBeenCalledWith(expect.objectContaining({ id: 'dep-new-002' }));
+  expect(store.addReceipt).toHaveBeenCalledWith(expect.objectContaining({ rerunOf: 'dep-old-001' }));
+  expect(lines.some(l => l.includes('dep-new-002'))).toBe(true);
+  expect(process.exitCode).toBeFalsy();
+});
+
+it('resolves a local name/alias to its deployment id before calling the API', async () => {
+  store.findDeployment.mockReturnValue({ id: 'dep-resolved-999' });
+  api.rerunDeployment.mockResolvedValue({ deployment_id: 'dep-new-003', workload_type: 'job', status: 'provisioning' });
+
+  await rerunCommand(config, ['my-job-alias'], chalk);
+
+  expect(api.rerunDeployment).toHaveBeenCalledWith(config, 'dep-resolved-999');
+});
+
+it('reports an error and sets exitCode on failure', async () => {
+  api.rerunDeployment.mockRejectedValue(new Error('deployment not found'));
+
+  await rerunCommand(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