diff --git a/packages/cli/src/cli/commands/fleet.test.ts b/packages/cli/src/cli/commands/fleet.test.ts index 19fa658d9..827423bf7 100644 --- a/packages/cli/src/cli/commands/fleet.test.ts +++ b/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; + 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 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 did not + // abort the command in enrollment mode. + expect(harness.env.RELAY_NODE_TOKEN).toBe('nt_secret'); + }); + + it('errors when neither 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 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('../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('../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 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(); diff --git a/packages/cli/src/cli/commands/fleet.ts b/packages/cli/src/cli/commands/fleet.ts index f9d17e6a2..a0d77a9fb 100644 --- a/packages/cli/src/cli/commands/fleet.ts +++ b/packages/cli/src/cli/commands/fleet.ts @@ -4,16 +4,18 @@ import { pathToFileURL } from 'node:url'; import type { Command } from 'commander'; import { createJiti } from 'jiti'; import { HarnessDriverClient } from '@agent-relay/harness-driver'; +import { enrollFleetNode, type FleetNodeEnrollment } from '@agent-relay/cloud'; import type { FleetNodeDefinition } from '@agent-relay/fleet'; // Namespace import sidesteps bun --compile's named-import validation against the // package .d.ts (see cli/lib/fleet-sidecar.ts). import * as fleetSdk from '@agent-relay/fleet'; const { isFleetNodeDefinition } = fleetSdk; -import { withDefaults, type CoreDependencies } from './core.js'; +import { withDefaults, type CoreDependencies, type CoreProjectPaths } from './core.js'; import { readBrokerConnection, startBrokerWithPortFallback } from '../lib/broker-lifecycle.js'; import { buildNodeSupervision, + createImplicitLocalFleetNode, fleetStatusPath, readFleetSidecarStatus, serveFleetSidecar, @@ -32,6 +34,7 @@ export interface FleetCommandDependencies { core: CoreDependencies; sdk: SdkCommandDeps; loadNodeDefinition: (file: string) => Promise; + enrollFleetNode: typeof enrollFleetNode; log: (...args: unknown[]) => void; warn: (...args: unknown[]) => void; error: (...args: unknown[]) => void; @@ -45,6 +48,7 @@ function withFleetDefaults(overrides: Partial = {}): F core, sdk, loadNodeDefinition, + enrollFleetNode, log: (...args: unknown[]) => console.log(...args), warn: (...args: unknown[]) => console.warn(...args), error: (...args: unknown[]) => console.error(...args), @@ -63,12 +67,20 @@ export function registerFleetCommands( group .command('serve') .description('Serve a fleet node definition') - .argument('', 'TS/JS node definition file') + .argument('[file]', 'TS/JS node definition file (optional when --enrollment-token is provided)') .option('--name ', 'Override node name') .option('--workspace ', 'Workspace key for broker registration and trigger sync') .option('--max-agents ', 'Override maximum managed agents for this node') .option('--base-url ', 'Override Relaycast API base URL') - .action(async (file: string, options: Record) => { + .option( + '--enrollment-token ', + 'One-time Cloud enrollment token (ocl_node_enr_...) to register this node' + ) + .option( + '--enrollment-url ', + 'Cloud enrollment endpoint that redeems the token (e.g. https://agentrelay.com/api/v1/fleet/register)' + ) + .action(async (file: string | undefined, options: Record) => { try { await runFleetServe(file, options, deps); } catch (error) { @@ -124,22 +136,134 @@ export async function loadNodeDefinition(file: string): Promise, + nameOption: string | undefined, + maxAgentsOverride: number | undefined, + deps: FleetCommandDependencies +): Promise { + const enrollmentToken = typeof options.enrollmentToken === 'string' ? options.enrollmentToken.trim() : ''; + const enrollmentUrl = typeof options.enrollmentUrl === 'string' ? options.enrollmentUrl.trim() : ''; + if (enrollmentUrl && !enrollmentToken) { + throw new Error('--enrollment-url requires --enrollment-token.'); + } + if (!enrollmentToken) { + return undefined; + } + + const enrollment = await deps.enrollFleetNode({ + enrollmentToken, + enrollmentUrl, + ...(nameOption ? { name: nameOption } : {}), + ...(maxAgentsOverride !== undefined ? { maxAgents: maxAgentsOverride } : {}), + }); + deps.core.env.RELAY_NODE_TOKEN = enrollment.nodeToken; + deps.core.env.RELAY_BASE_URL = enrollment.relaycastUrl; + deps.log( + `Enrolled fleet node "${enrollment.nodeName}"${ + enrollment.nodeId ? ` (${enrollment.nodeId})` : '' + } in workspace ${enrollment.relayWorkspaceId}.` + ); + return enrollment; +} + +/** + * Wires the explicit `--workspace`/`--base-url` overrides into the broker env so + * `startBrokerWithPortFallback` binds the node to the right workspace and origin. + * Returns the resolved values for downstream sidecar wiring. + * + * Precedence for the Relaycast origin: enrollment is the source of truth (it + * already wrote RELAY_BASE_URL during the exchange) and an explicit `--base-url` + * is the only thing that overrides it; the override must reach the broker via + * RELAY_BASE_URL, not just `serveFleetSidecar`. + */ +function applyServeEnvOverrides( + options: Record, + deps: FleetCommandDependencies +): { workspaceKey: string; baseUrlOverride: string } { + const workspaceKey = typeof options.workspace === 'string' ? options.workspace.trim() : ''; + if (workspaceKey) { + deps.core.env.RELAY_WORKSPACE_KEY = workspaceKey; + deps.core.env.RELAY_API_KEY = workspaceKey; + } + + const baseUrlOverride = typeof options.baseUrl === 'string' ? options.baseUrl.trim() : ''; + if (baseUrlOverride) { + deps.core.env.RELAY_BASE_URL = baseUrlOverride; + } + + return { workspaceKey, baseUrlOverride }; +} + +function createImplicitServeNodeDefinition(input: { + paths: CoreProjectPaths; + enrollment: FleetNodeEnrollment | undefined; + nameOption: string | undefined; + maxAgentsOverride: number | undefined; + deps: FleetCommandDependencies; +}): FleetNodeDefinition { + return createImplicitLocalFleetNode({ + paths: input.paths, + teamsConfig: input.deps.core.loadTeamsConfig(input.paths.projectRoot), + // Name precedence (kept identical to nameOverride in runFleetServe so the + // implicit node definition and the sidecar registration always agree): + // --name > enrollment record's nodeName > createImplicitLocalFleetNode's + // projectRoot-basename default. + name: input.nameOption ?? input.enrollment?.nodeName, + ...(input.maxAgentsOverride !== undefined ? { maxAgents: input.maxAgentsOverride } : {}), + }); +} + async function runFleetServe( - file: string, + file: string | undefined, options: Record, deps: FleetCommandDependencies ): Promise { - const nodeDefinition = await deps.loadNodeDefinition(file); const maxAgentsOverride = parsePositiveIntegerOption(options.maxAgents, '--max-agents'); + const nameOption = typeof options.name === 'string' ? options.name : undefined; const paths = deps.core.getProjectPaths(); deps.core.fs.mkdirSync(paths.dataDir, { recursive: true }); - const workspaceKey = typeof options.workspace === 'string' ? options.workspace.trim() : ''; - if (workspaceKey) { - deps.core.env.RELAY_WORKSPACE_KEY = workspaceKey; - deps.core.env.RELAY_API_KEY = workspaceKey; + // The `` node-def is OPTIONAL in enrollment mode — identity/name/ + // capabilities come from the enrollment record; a ``, when present, + // overrides/augments it. + // + // Load + validate the `` BEFORE redeeming the one-time enrollment token: + // the token is single-use, so a missing/invalid file must fail fast rather than + // burn the token on a run that can't succeed (the durable creds returned by the + // exchange would never be persisted, forcing the operator to mint a fresh token + // just to fix a local file error). + const fileDefinition = file ? await deps.loadNodeDefinition(file) : undefined; + + const enrollment = await maybeEnrollFleetNode(options, nameOption, maxAgentsOverride, deps); + if (!enrollment && !file) { + throw new Error('A node definition is required unless --enrollment-token is provided.'); } + const nodeDefinition = + fileDefinition ?? + createImplicitServeNodeDefinition({ + paths, + enrollment, + nameOption, + maxAgentsOverride, + deps, + }); + + const { workspaceKey, baseUrlOverride } = applyServeEnvOverrides(options, deps); + + // An enrolled node prefers the name from the enrollment record; a --name flag + // (already forwarded to the exchange) still wins through nameOption. + const nameOverride = nameOption ?? enrollment?.nodeName ?? undefined; + const baseUrl = baseUrlOverride || enrollment?.relaycastUrl || undefined; + const dashboardPort = Number.parseInt(deps.core.env.AGENT_RELAY_DASHBOARD_PORT ?? '3888', 10) || 3888; const started = await startBrokerWithPortFallback(paths, dashboardPort, deps.core); const connection = connectionFromFile(paths.dataDir); @@ -152,11 +276,17 @@ async function runFleetServe( definition: nodeDefinition, connection, workspaceKey: workspaceKey || started.relay.workspaceKey, - baseUrl: typeof options.baseUrl === 'string' ? options.baseUrl : undefined, - nameOverride: typeof options.name === 'string' ? options.name : undefined, + baseUrl, + nameOverride, maxAgentsOverride, supervision: buildNodeSupervision({ - argv: deps.core.argv, + // Strip the one-time --enrollment-token/--enrollment-url flags from the + // supervised argv. The broker restarts supervised sidecars by re-executing + // this argv; replaying a consumed enrollment token would fail the exchange + // and the node would never come back. The durable credentials minted by the + // first exchange live in the supervision env (RELAY_NODE_TOKEN / + // RELAY_BASE_URL), so restarts take the durable path with no flags to redeem. + argv: stripEnrollmentFlags(deps.core.argv), cwd: process.cwd(), env: deps.core.env, }), @@ -228,6 +358,33 @@ function connectionFromFile(dataDir: string): FleetBrokerConnection { return { url: conn.url, apiKey: conn.api_key }; } +/** + * Removes the one-time enrollment flags (and their values) from a captured argv + * so the broker's supervised-restart replay never re-redeems a consumed token. + * Handles both `--flag value` and `--flag=value` forms. The durable credentials + * minted by the first exchange are carried in the supervision env instead. + */ +export function stripEnrollmentFlags(argv: readonly string[]): string[] { + const oneTimeFlags = new Set(['--enrollment-token', '--enrollment-url']); + const result: string[] = []; + for (let i = 0; i < argv.length; i += 1) { + const arg = argv[i]!; + const eqIndex = arg.indexOf('='); + const flagName = eqIndex === -1 ? arg : arg.slice(0, eqIndex); + if (oneTimeFlags.has(flagName)) { + // `--flag=value` carries its value inline; `--flag value` consumes the next + // token. Skip the following token only when this flag had no inline `=value` + // and a value token actually follows. + if (eqIndex === -1 && i + 1 < argv.length) { + i += 1; + } + continue; + } + result.push(arg); + } + return result; +} + function parsePositiveIntegerOption(value: unknown, label: string): number | undefined { if (value === undefined) { return undefined; diff --git a/packages/cloud/src/fleet.test.ts b/packages/cloud/src/fleet.test.ts new file mode 100644 index 000000000..70affa0d5 --- /dev/null +++ b/packages/cloud/src/fleet.test.ts @@ -0,0 +1,179 @@ +import { describe, expect, it, vi } from 'vitest'; + +import { enrollFleetNode } from './fleet.js'; + +function jsonResponse(body: unknown, init: ResponseInit = {}): Response { + return new Response(JSON.stringify(body), { + ...init, + headers: { + 'content-type': 'application/json', + ...(init.headers ?? {}), + }, + }); +} + +const REGISTER_URL = 'https://agentrelay.com/api/v1/fleet/register'; + +describe('enrollFleetNode', () => { + it('exchanges a one-time token for node credentials', async () => { + const fetchImpl = vi.fn(async () => + jsonResponse({ + 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 result = await enrollFleetNode({ + enrollmentToken: ' ocl_node_enr_xyz ', + enrollmentUrl: REGISTER_URL, + name: 'kjglaptop', + maxAgents: 4, + capabilities: ['spawn:codex', ' spawn:codex ', ''], + tags: ['laptop'], + fetchImpl: fetchImpl as unknown as typeof fetch, + }); + + expect(fetchImpl).toHaveBeenCalledTimes(1); + const [calledUrl, calledInit] = fetchImpl.mock.calls[0] as [string, RequestInit]; + expect(calledUrl).toBe(REGISTER_URL); + expect(calledInit.method).toBe('POST'); + const sentBody = JSON.parse(String(calledInit.body)); + expect(sentBody).toMatchObject({ + enrollmentToken: 'ocl_node_enr_xyz', + name: 'kjglaptop', + maxAgents: 4, + capabilities: ['spawn:codex'], + tags: ['laptop'], + }); + expect(typeof sentBody.version).toBe('string'); + + expect(result).toMatchObject({ + nodeId: 'node_abc', + nodeName: 'kjglaptop', + nodeToken: 'nt_secret', + relayWorkspaceId: 'rw_123', + relaycastUrl: 'https://relaycast.example.com', + }); + }); + + it('derives the websocket url when the response omits it', async () => { + const fetchImpl = vi.fn(async () => + jsonResponse({ + nodeId: 'node_abc', + nodeName: 'n', + nodeToken: 'nt_secret', + relayWorkspaceId: 'rw_123', + relaycastUrl: 'https://relaycast.example.com', + }) + ); + + const result = await enrollFleetNode({ + enrollmentToken: 'ocl_node_enr_xyz', + enrollmentUrl: REGISTER_URL, + fetchImpl: fetchImpl as unknown as typeof fetch, + }); + + expect(result.websocketUrl).toBe('https://relaycast.example.com/v1/node/ws'); + }); + + it('throws a clear message when the token is expired/invalid/consumed (401)', async () => { + const fetchImpl = vi.fn(async () => jsonResponse({ error: 'Invalid enrollment token' }, { status: 401 })); + + await expect( + enrollFleetNode({ + enrollmentToken: 'ocl_node_enr_dead', + enrollmentUrl: REGISTER_URL, + fetchImpl: fetchImpl as unknown as typeof fetch, + }) + ).rejects.toThrow(/invalid, expired, or already used/i); + }); + + it('surfaces a rate-limit error on 429', async () => { + const fetchImpl = vi.fn(async () => jsonResponse({ error: 'Rate limit exceeded' }, { status: 429 })); + + await expect( + enrollFleetNode({ + enrollmentToken: 'ocl_node_enr_xyz', + enrollmentUrl: REGISTER_URL, + fetchImpl: fetchImpl as unknown as typeof fetch, + }) + ).rejects.toThrow(/rate limit/i); + }); + + it('does not dump HTML markup into the error when the URL is wrong (404)', async () => { + const html = `${'x'.repeat(500)}`; + const fetchImpl = vi.fn( + async () => new Response(html, { status: 404, headers: { 'content-type': 'text/html' } }) + ); + + await expect( + enrollFleetNode({ + enrollmentToken: 'ocl_node_enr_xyz', + enrollmentUrl: REGISTER_URL, + fetchImpl: fetchImpl as unknown as typeof fetch, + }) + ).rejects.toThrow(/Node enrollment failed: 404/); + }); + + it('rejects a response missing node credentials', async () => { + const fetchImpl = vi.fn(async () => jsonResponse({ nodeId: 'node_abc' })); + + await expect( + enrollFleetNode({ + enrollmentToken: 'ocl_node_enr_xyz', + enrollmentUrl: REGISTER_URL, + fetchImpl: fetchImpl as unknown as typeof fetch, + }) + ).rejects.toThrow(/missing node credentials/i); + }); + + it('rejects a response whose optional fields are the wrong type', async () => { + const fetchImpl = vi.fn(async () => + jsonResponse({ + nodeToken: 'nt_secret', + relayWorkspaceId: 'rw_123', + relaycastUrl: 'https://relaycast.example.com', + // The required trio is present, but a malformed optional field (number, + // not string) must still be caught by the tightened type guard rather + // than silently coerced. + nodeId: 42, + }) + ); + + await expect( + enrollFleetNode({ + enrollmentToken: 'ocl_node_enr_xyz', + enrollmentUrl: REGISTER_URL, + fetchImpl: fetchImpl as unknown as typeof fetch, + }) + ).rejects.toThrow(/missing node credentials/i); + }); + + it('rejects a response with an empty relayWorkspaceId', async () => { + const fetchImpl = vi.fn(async () => + jsonResponse({ + nodeToken: 'nt_secret', + relayWorkspaceId: ' ', + relaycastUrl: 'https://relaycast.example.com', + }) + ); + + await expect( + enrollFleetNode({ + enrollmentToken: 'ocl_node_enr_xyz', + enrollmentUrl: REGISTER_URL, + fetchImpl: fetchImpl as unknown as typeof fetch, + }) + ).rejects.toThrow(/missing node credentials/i); + }); + + it('requires a non-empty enrollment token', async () => { + await expect(enrollFleetNode({ enrollmentToken: ' ', enrollmentUrl: REGISTER_URL })).rejects.toThrow( + /enrollment token is required/i + ); + }); +}); diff --git a/packages/cloud/src/fleet.ts b/packages/cloud/src/fleet.ts new file mode 100644 index 000000000..3c08d1674 --- /dev/null +++ b/packages/cloud/src/fleet.ts @@ -0,0 +1,175 @@ +import os from 'node:os'; + +import { defaultApiUrl } from './types.js'; + +/** + * Credentials returned by the Cloud node-enrollment register endpoint + * (`/api/v1/fleet/register`). A one-time enrollment token is exchanged for these + * long-lived node credentials, which then configure the served fleet node. + * + * Mirrors the response shape of + * cloud/packages/web/app/api/v1/fleet/register/route.ts. + */ +export type FleetNodeEnrollment = { + nodeId: string; + nodeName: string; + nodeToken: string; + relayWorkspaceId: string; + relaycastUrl: string; + websocketUrl: string; +}; + +export type EnrollFleetNodeInput = { + /** One-time enrollment token minted by Cloud (`ocl_node_enr_...`). */ + enrollmentToken: string; + /** + * Cloud enrollment endpoint that redeems the token. Cloud mints this as + * `https:///api/v1/fleet/register`. + */ + enrollmentUrl: string; + /** Optional node name override; otherwise the enrollment record's name is used. */ + name?: string; + /** Optional capability override/augment for the registered node. */ + capabilities?: string[]; + /** Optional max-agents override for the registered node. */ + maxAgents?: number; + /** Optional tags override for the registered node. */ + tags?: string[]; + /** Optional version string reported to Cloud (defaults to host info). */ + version?: string; + fetchImpl?: typeof fetch; + signal?: AbortSignal; +}; + +function ensurePlainObject(value: unknown): value is Record { + return value !== null && typeof value === 'object' && !Array.isArray(value); +} + +function normalizeEnrollmentUrl(value?: string): string { + const raw = value?.trim(); + if (!raw) { + // Fall back to the configured Cloud API origin if the caller omitted a URL. + return `${defaultApiUrl().replace(/\/+$/, '')}/api/v1/fleet/register`; + } + try { + return new URL(raw).toString(); + } catch { + throw new Error(`Invalid enrollment URL: ${raw}`); + } +} + +function normalizeStringList(values?: string[]): string[] | undefined { + if (values === undefined) return undefined; + return [...new Set(values.map((value) => value.trim()).filter(Boolean))]; +} + +async function readJsonResponse(response: Response): Promise { + const contentType = response.headers.get('content-type') ?? ''; + if (contentType.includes('application/json')) { + return response.json().catch(() => null); + } + const text = await response.text().catch(() => ''); + return text ? { error: text } : null; +} + +function enrollmentError(response: Response, payload: unknown): Error { + // Prefer the JSON {error} field. For non-JSON bodies (e.g. an HTML 404 page + // when the URL is wrong) fall back to the status line rather than dumping the + // raw markup into the operator's terminal. + const rawError = + ensurePlainObject(payload) && typeof payload.error === 'string' ? payload.error.trim() : ''; + const looksLikeMarkup = rawError.startsWith('<') || rawError.length > 200; + const detail = rawError && !looksLikeMarkup ? rawError : `${response.status} ${response.statusText}`.trim(); + + // The register endpoint returns 401 "Invalid enrollment token" for tokens that + // are expired, already consumed, or never minted. Give the operator a clear, + // actionable message — enrollment tokens are one-time and short-lived. + if (response.status === 401 || /invalid enrollment token/i.test(detail)) { + return new Error( + 'Enrollment token is invalid, expired, or already used. Mint a fresh token from the Cloud "Enroll node" command and retry.' + ); + } + if (response.status === 429) { + return new Error('Enrollment rate limit exceeded; wait a moment and retry.'); + } + return new Error(`Node enrollment failed: ${detail}`); +} + +/** + * Validate the register endpoint's response carries the credentials we depend + * on. The required trio (nodeToken/relaycastUrl/relayWorkspaceId) must be present + * non-empty strings; the optional fields (nodeId/nodeName/websocketUrl), which we + * fill in with derived defaults below, must be strings when present so a + * malformed server response is rejected rather than silently coerced. + */ +function isFleetNodeEnrollment(value: unknown): value is Partial & { + nodeToken: string; + relaycastUrl: string; + relayWorkspaceId: string; +} { + if (!ensurePlainObject(value)) return false; + return ( + typeof value.nodeToken === 'string' && + value.nodeToken.trim().length > 0 && + typeof value.relaycastUrl === 'string' && + value.relaycastUrl.trim().length > 0 && + typeof value.relayWorkspaceId === 'string' && + value.relayWorkspaceId.trim().length > 0 && + (value.nodeId === undefined || typeof value.nodeId === 'string') && + (value.nodeName === undefined || typeof value.nodeName === 'string') && + (value.websocketUrl === undefined || typeof value.websocketUrl === 'string') + ); +} + +/** + * Exchange a one-time fleet-node enrollment token for durable node credentials. + * + * Models the worker enrollment exchange (registerCloudWorker in ./worker.ts): + * a single unauthenticated POST that redeems the token and returns credentials + * used to serve the node. Unlike worker registration there is no local store — + * an enrolled node is configured per-invocation from the returned credentials. + */ +export async function enrollFleetNode(input: EnrollFleetNodeInput): Promise { + const fetcher = input.fetchImpl ?? fetch; + const enrollmentToken = input.enrollmentToken.trim(); + if (!enrollmentToken) { + throw new Error('An enrollment token is required to enroll a fleet node.'); + } + const url = normalizeEnrollmentUrl(input.enrollmentUrl); + const capabilities = normalizeStringList(input.capabilities); + const tags = normalizeStringList(input.tags); + const name = input.name?.trim(); + + const response = await fetcher(url, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ + enrollmentToken, + ...(name ? { name } : {}), + ...(capabilities !== undefined ? { capabilities } : {}), + ...(input.maxAgents !== undefined ? { maxAgents: input.maxAgents } : {}), + ...(tags !== undefined ? { tags } : {}), + version: input.version?.trim() || `relay-cli/${os.platform()}-${os.arch()}`, + }), + signal: input.signal, + }); + + const payload = await readJsonResponse(response); + if (!response.ok) { + throw enrollmentError(response, payload); + } + if (!isFleetNodeEnrollment(payload)) { + throw new Error('Node enrollment response is missing node credentials.'); + } + return { + nodeId: typeof payload.nodeId === 'string' ? payload.nodeId : '', + nodeName: typeof payload.nodeName === 'string' ? payload.nodeName : (name ?? ''), + nodeToken: payload.nodeToken.trim(), + relayWorkspaceId: payload.relayWorkspaceId, + relaycastUrl: payload.relaycastUrl.replace(/\/+$/, ''), + websocketUrl: + typeof payload.websocketUrl === 'string' && payload.websocketUrl.trim() + ? payload.websocketUrl.trim() + : `${payload.relaycastUrl.replace(/\/+$/, '')}/v1/node/ws`, + }; +} diff --git a/packages/cloud/src/index.ts b/packages/cloud/src/index.ts index 5ca84fb41..3002a4531 100644 --- a/packages/cloud/src/index.ts +++ b/packages/cloud/src/index.ts @@ -67,6 +67,8 @@ export { type WorkerWorkflowRef, } from './worker.js'; +export { enrollFleetNode, type EnrollFleetNodeInput, type FleetNodeEnrollment } from './fleet.js'; + export { activeWorkspaceKey, readWorkspaceStore, diff --git a/tests/e2e/fleet/fleet-e2e.test.ts b/tests/e2e/fleet/fleet-e2e.test.ts index d557434a1..afcabc2a7 100644 --- a/tests/e2e/fleet/fleet-e2e.test.ts +++ b/tests/e2e/fleet/fleet-e2e.test.ts @@ -334,40 +334,57 @@ describe.skipIf(!pre.ok)('two-node fleet scenario matrix', () => { expect(scheduled.body.data.handler_node_id).toBe('node_b'); // least-loaded }, 60_000); - it('resume: a resumable spawn re-binds to the agent ORIGIN node (not an arbitrary target)', async () => { - const sessionRef = 'sess-resume-1'; - // First spawn is UNTARGETED → the engine picks the origin node by placement. - // We capture wherever it actually landed so the resume target is derived from - // the agent's real origin, not hard-coded (resume = targeted-origin spawn; - // the engine records origin_node_id but does not auto-route from session_ref). - const first = await invokeAction(engine, driverToken, 'spawn', { - cli: 'pool', - name: 'resumable-1', - session_ref: sessionRef, - }); - const originId = first.body.data.handler_node_id as string; // engine-chosen origin - const originName = originId === 'node_a' ? 'node-a' : 'node-b'; - const firstDone = await waitFor( - async () => { - const inv = await getInvocation(engine, driverToken, 'spawn', first.invocationId!); - return inv.status === 'completed' || inv.status === 'failed' ? inv : null; - }, - { label: 'resumable spawn settled', timeoutMs: 20_000 } - ); - expect(firstDone.status).toBe('completed'); // resumable spawn carried session_ref through token authority - - // Release, then resume the SAME session targeted at the recorded origin. - expect(await releaseAgent(engine, workspaceKey, 'resumable-1')).toBeLessThan(300); - const resume = await invokeAction(engine, driverToken, 'spawn', { - cli: 'pool', - name: 'resumable-1', - target_node: originName, - session_ref: sessionRef, - }); - expect(resume.status).toBe(201); - expect(resume.body.data.handler_node_id).toBe(originId); // resumed on the agent's origin node - expect(resume.body.data.dispatched_node_id).toBe(originId); - }, 45_000); + // This is the 7th scenario in the serial chain — by now both nodes are running + // several stub PTY children from the earlier spawn scenarios, so the broker + + // sidecar are under real contention and the FIRST (untargeted) spawn's settle + // can occasionally exceed a tight deadline (observed `last=null` ⇒ the + // invocation simply hadn't reached a terminal status yet, not a logic fault). + // The origin-rebind correctness (the actual subject of this test, asserted on + // the resume response below) is unaffected — so we give the settle a realistic + // deadline and a bounded retry rather than weakening any assertion. The retry + // re-runs the whole body, so we first release any `resumable-1` left bound by a + // prior timed-out attempt (the release at the end is skipped when settle throws) + // to keep each attempt starting from a clean slate. + it( + 'resume: a resumable spawn re-binds to the agent ORIGIN node (not an arbitrary target)', + { retry: 2 }, + async () => { + const sessionRef = 'sess-resume-1'; + await releaseAgent(engine, workspaceKey, 'resumable-1'); // idempotent cleanup for retries + // First spawn is UNTARGETED → the engine picks the origin node by placement. + // We capture wherever it actually landed so the resume target is derived from + // the agent's real origin, not hard-coded (resume = targeted-origin spawn; + // the engine records origin_node_id but does not auto-route from session_ref). + const first = await invokeAction(engine, driverToken, 'spawn', { + cli: 'pool', + name: 'resumable-1', + session_ref: sessionRef, + }); + const originId = first.body.data.handler_node_id as string; // engine-chosen origin + const originName = originId === 'node_a' ? 'node-a' : 'node-b'; + const firstDone = await waitFor( + async () => { + const inv = await getInvocation(engine, driverToken, 'spawn', first.invocationId!); + return inv.status === 'completed' || inv.status === 'failed' ? inv : null; + }, + { label: 'resumable spawn settled', timeoutMs: 30_000, intervalMs: 300 } + ); + expect(firstDone.status).toBe('completed'); // resumable spawn carried session_ref through token authority + + // Release, then resume the SAME session targeted at the recorded origin. + expect(await releaseAgent(engine, workspaceKey, 'resumable-1')).toBeLessThan(300); + const resume = await invokeAction(engine, driverToken, 'spawn', { + cli: 'pool', + name: 'resumable-1', + target_node: originName, + session_ref: sessionRef, + }); + expect(resume.status).toBe(201); + expect(resume.body.data.handler_node_id).toBe(originId); // resumed on the agent's origin node + expect(resume.body.data.dispatched_node_id).toBe(originId); + }, + 60_000 + ); it('placement failure: spawning a capability no targeted node advertises fails with capability_mismatch', async () => { const res = await invokeAction(engine, driverToken, 'spawn', {