feat(cli): support fleet-node enrollment in fleet serve

packages/cli/src/cli/commands/fleet.test.ts

@@ -1,12 +1,13 @@
 import { once } from 'node:events';
 import path from 'node:path';
 
+import { Command } from 'commander';
 import { WebSocketServer } from 'ws';
 import { describe, expect, it, vi } from 'vitest';
 import { z } from 'zod';
 import { action, defineNode, onMessage } from '@agent-relay/fleet';
 
-import { loadNodeDefinition } from './fleet.js';
+import { loadNodeDefinition, registerFleetCommands, stripEnrollmentFlags } from './fleet.js';
 import { startFleetSidecar, serveFleetSidecar } from '../lib/fleet-sidecar.js';
 
 describe('fleet command support', () => {
@@ -142,6 +143,341 @@ describe('fleet command support', () => {
     broker.close();
   });
 
+  describe('fleet serve enrollment flags', () => {
+    function buildServeHarness(
+      overrides: {
+        enrollFleetNode?: ReturnType<typeof vi.fn>;
+        env?: NodeJS.ProcessEnv;
+      } = {}
+    ) {
+      const env: NodeJS.ProcessEnv = overrides.env ?? {};
+      const errors: string[] = [];
+      const exit = vi.fn(() => {
+        throw new Error('__exit__');
+      });
+      // Stop the flow right after enrollment by failing broker startup with a
+      // sentinel, so the test exercises flag parsing + the token exchange only.
+      const createRelay = vi.fn(() => {
+        throw new Error('__stop_after_enrollment__');
+      });
+      const enroll =
+        overrides.enrollFleetNode ??
+        vi.fn(async () => ({
+          nodeId: 'node_abc',
+          nodeName: 'kjglaptop',
+          nodeToken: 'nt_secret',
+          relayWorkspaceId: 'rw_123',
+          relaycastUrl: 'https://relaycast.example.com',
+          websocketUrl: 'https://relaycast.example.com/v1/node/ws',
+        }));
+
+      const core = {
+        getProjectPaths: () => ({ projectRoot: '/tmp/proj', dataDir: '/tmp/proj/.data' }),
+        loadTeamsConfig: () => null,
+        createRelay,
+        fs: { mkdirSync: vi.fn() },
+        env,
+        argv: ['node', 'agent-relay'],
+        onSignal: vi.fn(),
+        isPortInUse: vi.fn(async () => false),
+        exit,
+      } as never;
+
+      const program = new Command();
+      program.exitOverride();
+      registerFleetCommands(program, {
+        core,
+        enrollFleetNode: enroll as never,
+        error: (...args: unknown[]) => errors.push(args.join(' ')),
+        log: () => undefined,
+        warn: () => undefined,
+        exit: exit as never,
+      });
+
+      return { program, enroll, env, errors, exit };
+    }
+
+    it('accepts --enrollment-token/--enrollment-url and exchanges the token', async () => {
+      const harness = buildServeHarness();
+
+      await harness.program
+        .parseAsync(
+          [
+            'fleet',
+            'serve',
+            '--enrollment-token',
+            'ocl_node_enr_xyz',
+            '--enrollment-url',
+            'https://agentrelay.com/api/v1/fleet/register',
+            '--name',
+            'kjglaptop',
+            '--max-agents',
+            '4',
+          ],
+          { from: 'user' }
+        )
+        .catch(() => undefined);
+
+      expect(harness.enroll).toHaveBeenCalledTimes(1);
+      expect(harness.enroll).toHaveBeenCalledWith(
+        expect.objectContaining({
+          enrollmentToken: 'ocl_node_enr_xyz',
+          enrollmentUrl: 'https://agentrelay.com/api/v1/fleet/register',
+          name: 'kjglaptop',
+          maxAgents: 4,
+        })
+      );
+      // The exchange result is wired into the broker env before serving.
+      expect(harness.env.RELAY_NODE_TOKEN).toBe('nt_secret');
+      expect(harness.env.RELAY_BASE_URL).toBe('https://relaycast.example.com');
+    });
+
+    it('lets an explicit --base-url override the enrollment relaycast url in the broker env', async () => {
+      const harness = buildServeHarness();
+
+      await harness.program
+        .parseAsync(
+          [
+            'fleet',
+            'serve',
+            '--enrollment-token',
+            'ocl_node_enr_xyz',
+            '--base-url',
+            'https://override.example.com',
+          ],
+          { from: 'user' }
+        )
+        .catch(() => undefined);
+
+      // Enrollment first writes its relaycastUrl, then the explicit --base-url
+      // overrides it so the broker (started from the env) binds to the override.
+      expect(harness.env.RELAY_BASE_URL).toBe('https://override.example.com');
+    });
+
+    it('keeps the enrollment relaycast url in the broker env when --base-url is omitted', async () => {
+      const harness = buildServeHarness();
+
+      await harness.program
+        .parseAsync(['fleet', 'serve', '--enrollment-token', 'ocl_node_enr_xyz'], { from: 'user' })
+        .catch(() => undefined);
+
+      expect(harness.env.RELAY_BASE_URL).toBe('https://relaycast.example.com');
+    });
+
+    it('serves an enrolled node without a <file> argument', async () => {
+      const harness = buildServeHarness();
+
+      await harness.program
+        .parseAsync(['fleet', 'serve', '--enrollment-token', 'ocl_node_enr_xyz'], { from: 'user' })
+        .catch(() => undefined);
+
+      expect(harness.enroll).toHaveBeenCalledTimes(1);
+      // Reaching broker startup (the sentinel) proves the missing <file> did not
+      // abort the command in enrollment mode.
+      expect(harness.env.RELAY_NODE_TOKEN).toBe('nt_secret');
+    });
+
+    it('errors when neither <file> nor --enrollment-token is provided', async () => {
+      const harness = buildServeHarness();
+
+      await harness.program.parseAsync(['fleet', 'serve'], { from: 'user' }).catch(() => undefined);
+
+      expect(harness.enroll).not.toHaveBeenCalled();
+      expect(harness.errors.join('\n')).toMatch(/node definition <file> is required/i);
+    });
+
+    it('rejects --enrollment-url without --enrollment-token', async () => {
+      const harness = buildServeHarness();
+
+      await harness.program
+        .parseAsync(['fleet', 'serve', '--enrollment-url', 'https://agentrelay.com/api/v1/fleet/register'], {
+          from: 'user',
+        })
+        .catch(() => undefined);
+
+      expect(harness.enroll).not.toHaveBeenCalled();
+      expect(harness.errors.join('\n')).toMatch(/--enrollment-url requires --enrollment-token/i);
+    });
+
+    it('prefers --name over the enrollment nodeName when building the implicit node', async () => {
+      vi.resetModules();
+
+      const createImplicitLocalFleetNode = vi.fn(() => defineNode({ name: 'placeholder', capabilities: {} }));
+      vi.doMock('../lib/fleet-sidecar.js', async () => {
+        const actual =
+          await vi.importActual<typeof import('../lib/fleet-sidecar.js')>('../lib/fleet-sidecar.js');
+        return { ...actual, createImplicitLocalFleetNode };
+      });
+
+      const { registerFleetCommands: registerWithMock } = await import('./fleet.js');
+
+      const enroll = vi.fn(async () => ({
+        nodeId: 'node_abc',
+        nodeName: 'enrollment-name',
+        nodeToken: 'nt_secret',
+        relayWorkspaceId: 'rw_123',
+        relaycastUrl: 'https://relaycast.example.com',
+        websocketUrl: 'https://relaycast.example.com/v1/node/ws',
+      }));
+      const core = {
+        getProjectPaths: () => ({ projectRoot: '/tmp/proj', dataDir: '/tmp/proj/.data' }),
+        loadTeamsConfig: () => null,
+        createRelay: vi.fn(() => {
+          throw new Error('__stop_after_enrollment__');
+        }),
+        fs: { mkdirSync: vi.fn() },
+        env: {} as NodeJS.ProcessEnv,
+        argv: ['node', 'agent-relay'],
+        onSignal: vi.fn(),
+        isPortInUse: vi.fn(async () => false),
+        exit: vi.fn(() => {
+          throw new Error('__exit__');
+        }),
+      } as never;
+
+      const program = new Command();
+      program.exitOverride();
+      registerWithMock(program, {
+        core,
+        enrollFleetNode: enroll as never,
+        error: () => undefined,
+        log: () => undefined,
+        warn: () => undefined,
+        exit: (() => {
+          throw new Error('__exit__');
+        }) as never,
+      });
+
+      await program
+        .parseAsync(['fleet', 'serve', '--enrollment-token', 'ocl_node_enr_xyz', '--name', 'cli-name'], {
+          from: 'user',
+        })
+        .catch(() => undefined);
+
+      expect(createImplicitLocalFleetNode).toHaveBeenCalledWith(
+        expect.objectContaining({ name: 'cli-name' })
+      );
+
+      vi.doUnmock('../lib/fleet-sidecar.js');
+      vi.resetModules();
+    });
+
+    it('falls back to the enrollment nodeName when --name is omitted', async () => {
+      vi.resetModules();
+
+      const createImplicitLocalFleetNode = vi.fn(() => defineNode({ name: 'placeholder', capabilities: {} }));
+      vi.doMock('../lib/fleet-sidecar.js', async () => {
+        const actual =
+          await vi.importActual<typeof import('../lib/fleet-sidecar.js')>('../lib/fleet-sidecar.js');
+        return { ...actual, createImplicitLocalFleetNode };
+      });
+
+      const { registerFleetCommands: registerWithMock } = await import('./fleet.js');
+
+      const enroll = vi.fn(async () => ({
+        nodeId: 'node_abc',
+        nodeName: 'enrollment-name',
+        nodeToken: 'nt_secret',
+        relayWorkspaceId: 'rw_123',
+        relaycastUrl: 'https://relaycast.example.com',
+        websocketUrl: 'https://relaycast.example.com/v1/node/ws',
+      }));
+      const core = {
+        getProjectPaths: () => ({ projectRoot: '/tmp/proj', dataDir: '/tmp/proj/.data' }),
+        loadTeamsConfig: () => null,
+        createRelay: vi.fn(() => {
+          throw new Error('__stop_after_enrollment__');
+        }),
+        fs: { mkdirSync: vi.fn() },
+        env: {} as NodeJS.ProcessEnv,
+        argv: ['node', 'agent-relay'],
+        onSignal: vi.fn(),
+        isPortInUse: vi.fn(async () => false),
+        exit: vi.fn(() => {
+          throw new Error('__exit__');
+        }),
+      } as never;
+
+      const program = new Command();
+      program.exitOverride();
+      registerWithMock(program, {
+        core,
+        enrollFleetNode: enroll as never,
+        error: () => undefined,
+        log: () => undefined,
+        warn: () => undefined,
+        exit: (() => {
+          throw new Error('__exit__');
+        }) as never,
+      });
+
+      await program
+        .parseAsync(['fleet', 'serve', '--enrollment-token', 'ocl_node_enr_xyz'], { from: 'user' })
+        .catch(() => undefined);
+
+      expect(createImplicitLocalFleetNode).toHaveBeenCalledWith(
+        expect.objectContaining({ name: 'enrollment-name' })
+      );
+
+      vi.doUnmock('../lib/fleet-sidecar.js');
+      vi.resetModules();
+    });
+
+    it('validates the <file> before redeeming the one-time enrollment token', async () => {
+      const harness = buildServeHarness();
+
+      // A nonexistent node file must fail-fast WITHOUT burning the single-use
+      // enrollment token, so the operator can fix the path and retry the token.
+      await harness.program
+        .parseAsync(
+          ['fleet', 'serve', '/tmp/does-not-exist-node-def.ts', '--enrollment-token', 'ocl_node_enr_xyz'],
+          { from: 'user' }
+        )
+        .catch(() => undefined);
+
+      expect(harness.enroll).not.toHaveBeenCalled();
+      expect(harness.env.RELAY_NODE_TOKEN).toBeUndefined();
+    });
+  });
+
+  describe('stripEnrollmentFlags', () => {
+    it('removes --enrollment-token/--enrollment-url and their space-separated values', () => {
+      const argv = [
+        'node',
+        'agent-relay',
+        'fleet',
+        'serve',
+        '--enrollment-token',
+        'ocl_node_enr_xyz',
+        '--enrollment-url',
+        'https://agentrelay.com/api/v1/fleet/register',
+        '--name',
+        'kjglaptop',
+      ];
+
+      expect(stripEnrollmentFlags(argv)).toEqual([
+        'node',
+        'agent-relay',
+        'fleet',
+        'serve',
+        '--name',
+        'kjglaptop',
+      ]);
+    });
+
+    it('removes the --flag=value inline form without dropping the following token', () => {
+      const argv = ['fleet', 'serve', '--enrollment-token=ocl_node_enr_xyz', '--name', 'kjglaptop'];
+
+      expect(stripEnrollmentFlags(argv)).toEqual(['fleet', 'serve', '--name', 'kjglaptop']);
+    });
+
+    it('leaves argv untouched when no enrollment flags are present', () => {
+      const argv = ['fleet', 'serve', 'node.ts', '--base-url', 'https://relaycast.example.com'];
+
+      expect(stripEnrollmentFlags(argv)).toEqual(argv);
+    });
+  });
+
   it('keeps trigger sync idempotent across repeated node registrations', async () => {
     vi.resetModules();