feat: add minimal companion client CLI command

2026-02-18 11:02:21 -08:00
parent 7ca5d5bff5
commit 59c1033da0
6 changed files with 366 additions and 2 deletions
@@ -0,0 +1,152 @@
+import { Command } from 'commander';
+import { beforeEach, describe, expect, it, vi } from 'vitest';
+
+const {
+  mockLoadConfigSafe,
+  mockGetConfigPath,
+  mockRuntimeCtorArgs,
+  mockRuntimeInstances,
+} = vi.hoisted(() => {
+  const runtimeCtorArgs: Array<{ url: string; token?: string }> = [];
+  const runtimeInstances: Array<{
+    connect: ReturnType<typeof vi.fn>;
+    registerNode: ReturnType<typeof vi.fn>;
+    setNodeStatus: ReturnType<typeof vi.fn>;
+    subscribeAgentStream: ReturnType<typeof vi.fn>;
+    subscribeAgentTyping: ReturnType<typeof vi.fn>;
+    disconnect: ReturnType<typeof vi.fn>;
+  }> = [];
+
+  const loadConfigSafe = vi.fn(() => ({
+    config: {
+      server: {
+        port: 18888,
+        token: 'config-token',
+      },
+    },
+  }));
+
+  const getConfigPath = vi.fn(() => '/tmp/flynn-config.yaml');
+
+  return {
+    mockLoadConfigSafe: loadConfigSafe,
+    mockGetConfigPath: getConfigPath,
+    mockRuntimeCtorArgs: runtimeCtorArgs,
+    mockRuntimeInstances: runtimeInstances,
+  };
+});
+
+vi.mock('./shared.js', () => ({
+  loadConfigSafe: mockLoadConfigSafe,
+  getConfigPath: mockGetConfigPath,
+}));
+
+vi.mock('../companion/index.js', () => ({
+  CompanionRuntimeClient: class {
+    connect = vi.fn(async () => undefined);
+    registerNode = vi.fn(async ({ nodeId, role, capabilities }: { nodeId: string; role: string; capabilities: string[] }) => ({
+      registered: true,
+      node: { id: nodeId, role },
+      protocol: { serverVersion: 1, clientVersion: 1, negotiatedVersion: 1 },
+      capabilities: { declared: capabilities, enabled: capabilities },
+    }));
+    setNodeStatus = vi.fn(async () => ({ updated: true, node: { id: 'n', role: 'companion' } }));
+    subscribeAgentStream = vi.fn(() => () => undefined);
+    subscribeAgentTyping = vi.fn(() => () => undefined);
+    disconnect = vi.fn(() => undefined);
+
+    constructor(opts: { url: string; token?: string }) {
+      mockRuntimeCtorArgs.push(opts);
+      mockRuntimeInstances.push(this);
+    }
+  },
+}));
+
+describe('companion command', () => {
+  beforeEach(() => {
+    vi.clearAllMocks();
+    mockRuntimeCtorArgs.length = 0;
+    mockRuntimeInstances.length = 0;
+    mockLoadConfigSafe.mockReturnValue({
+      config: {
+        server: {
+          port: 18888,
+          token: 'config-token',
+        },
+      },
+    });
+    mockGetConfigPath.mockReturnValue('/tmp/flynn-config.yaml');
+    process.exitCode = undefined;
+  });
+
+  it('uses config-derived gateway url/token by default', async () => {
+    const logSpy = vi.spyOn(console, 'log').mockImplementation(() => undefined);
+    const errSpy = vi.spyOn(console, 'error').mockImplementation(() => undefined);
+    const program = new Command();
+    const { registerCompanionCommand } = await import('./companion.js');
+    registerCompanionCommand(program);
+
+    await program.parseAsync(['node', 'test', 'companion', '--once']);
+
+    expect(mockGetConfigPath).toHaveBeenCalledOnce();
+    expect(mockRuntimeCtorArgs).toEqual([{ url: 'ws://127.0.0.1:18888', token: 'config-token' }]);
+    expect(mockRuntimeInstances[0]?.connect).toHaveBeenCalledOnce();
+    expect(mockRuntimeInstances[0]?.registerNode).toHaveBeenCalledOnce();
+    expect(mockRuntimeInstances[0]?.setNodeStatus).toHaveBeenCalledOnce();
+    expect(mockRuntimeInstances[0]?.disconnect).toHaveBeenCalled();
+    expect(errSpy).not.toHaveBeenCalled();
+    expect(process.exitCode).toBeUndefined();
+
+    logSpy.mockRestore();
+    errSpy.mockRestore();
+  });
+
+  it('prefers explicit url/token and capability overrides', async () => {
+    const logSpy = vi.spyOn(console, 'log').mockImplementation(() => undefined);
+    const errSpy = vi.spyOn(console, 'error').mockImplementation(() => undefined);
+    const program = new Command();
+    const { registerCompanionCommand } = await import('./companion.js');
+    registerCompanionCommand(program);
+
+    await program.parseAsync([
+      'node',
+      'test',
+      'companion',
+      '--once',
+      '--url',
+      'ws://10.0.0.5:19000',
+      '--token',
+      'override-token',
+      '--node-id',
+      'test-node',
+      '--capability',
+      'ui.canvas',
+      'node.push.register',
+    ]);
+
+    expect(mockRuntimeCtorArgs).toEqual([{ url: 'ws://10.0.0.5:19000', token: 'override-token' }]);
+    expect(mockRuntimeInstances[0]?.registerNode).toHaveBeenCalledWith(expect.objectContaining({
+      nodeId: 'test-node',
+      capabilities: ['ui.canvas', 'node.push.register'],
+    }));
+    expect(errSpy).not.toHaveBeenCalled();
+
+    logSpy.mockRestore();
+    errSpy.mockRestore();
+  });
+
+  it('sets process exit code when options are invalid', async () => {
+    const errSpy = vi.spyOn(console, 'error').mockImplementation(() => undefined);
+    const program = new Command();
+    const { registerCompanionCommand } = await import('./companion.js');
+    registerCompanionCommand(program);
+
+    await program.parseAsync(['node', 'test', 'companion', '--once', '--heartbeat', '0']);
+
+    expect(errSpy).toHaveBeenCalled();
+    expect(process.exitCode).toBe(1);
+
+    errSpy.mockRestore();
+  });
+});
+
@@ -0,0 +1,189 @@
+import type { Command } from 'commander';
+import { hostname } from 'node:os';
+import { randomUUID } from 'node:crypto';
+import { CompanionRuntimeClient } from '../companion/index.js';
+import type { SetNodeStatusInput } from '../companion/index.js';
+import { getConfigPath, loadConfigSafe } from './shared.js';
+
+type CompanionPlatform = SetNodeStatusInput['platform'];
+
+interface CompanionCommandOptions {
+  config?: string;
+  url?: string;
+  token?: string;
+  nodeId?: string;
+  role?: string;
+  platform?: CompanionPlatform;
+  capability?: string[];
+  heartbeat?: string;
+  once?: boolean;
+}
+
+function resolveGatewayUrl(options: CompanionCommandOptions, configPath: string): string {
+  if (options.url && options.url.trim().length > 0) {
+    return options.url.trim();
+  }
+
+  const loaded = loadConfigSafe(configPath);
+  if (loaded.config) {
+    const port = loaded.config.server.port;
+    return `ws://127.0.0.1:${port}`;
+  }
+
+  return 'ws://127.0.0.1:18800';
+}
+
+function resolveGatewayToken(options: CompanionCommandOptions, configPath: string): string | undefined {
+  if (options.token && options.token.trim().length > 0) {
+    return options.token.trim();
+  }
+
+  const loaded = loadConfigSafe(configPath);
+  return loaded.config?.server.token;
+}
+
+function resolveCapabilities(platform: CompanionPlatform, provided?: string[]): string[] {
+  if (provided && provided.length > 0) {
+    return provided.map((v) => v.trim()).filter(Boolean);
+  }
+
+  if (platform === 'ios' || platform === 'macos') {
+    return ['ui.canvas', 'node.status.write', 'node.location.write', 'node.push.register'];
+  }
+  if (platform === 'android') {
+    return ['ui.canvas', 'node.status.write', 'node.location.write', 'node.push.register'];
+  }
+  return ['ui.canvas', 'node.status.write'];
+}
+
+function resolveNodeId(options: CompanionCommandOptions, platform: CompanionPlatform): string {
+  if (options.nodeId && options.nodeId.trim().length > 0) {
+    return options.nodeId.trim();
+  }
+  return `${platform}-${hostname()}-${randomUUID().slice(0, 8)}`;
+}
+
+function parseHeartbeatSeconds(value: string | undefined): number {
+  const raw = value ?? '30';
+  const parsed = Number.parseInt(raw, 10);
+  if (!Number.isFinite(parsed) || parsed <= 0 || parsed > 86_400) {
+    throw new Error('heartbeat must be an integer between 1 and 86400 seconds');
+  }
+  return parsed;
+}
+
+async function publishHeartbeat(
+  runtime: CompanionRuntimeClient,
+  platform: CompanionPlatform,
+): Promise<void> {
+  await runtime.setNodeStatus({
+    platform,
+    statusText: 'heartbeat',
+    powerSource: 'unknown',
+  });
+}
+
+export async function runCompanionSession(options: CompanionCommandOptions): Promise<void> {
+  const configPath = options.config ?? getConfigPath();
+  const platform: CompanionPlatform = options.platform ?? 'macos';
+  const gatewayUrl = resolveGatewayUrl(options, configPath);
+  const gatewayToken = resolveGatewayToken(options, configPath);
+  const role = options.role?.trim() || 'companion';
+  const nodeId = resolveNodeId(options, platform);
+  const capabilities = resolveCapabilities(platform, options.capability);
+  const heartbeatSeconds = parseHeartbeatSeconds(options.heartbeat);
+
+  const runtime = new CompanionRuntimeClient({
+    url: gatewayUrl,
+    token: gatewayToken,
+  });
+
+  const stopSignals: NodeJS.Signals[] = ['SIGINT', 'SIGTERM'];
+  let heartbeatTimer: NodeJS.Timeout | null = null;
+
+  const cleanup = (): void => {
+    if (heartbeatTimer) {
+      clearInterval(heartbeatTimer);
+      heartbeatTimer = null;
+    }
+    runtime.disconnect(1000, 'Companion shutting down');
+  };
+
+  for (const signal of stopSignals) {
+    process.once(signal, cleanup);
+  }
+
+  runtime.subscribeAgentStream((data) => {
+    const payload = data as { sessionId?: string; content?: string };
+    const session = payload.sessionId ? ` (${payload.sessionId})` : '';
+    if (payload.content) {
+      console.log(`[agent.stream${session}] ${payload.content}`);
+    }
+  });
+
+  runtime.subscribeAgentTyping((data) => {
+    const payload = data as { sessionId?: string; phase?: string };
+    const session = payload.sessionId ? ` (${payload.sessionId})` : '';
+    const phase = payload.phase ?? 'typing';
+    console.log(`[agent.typing${session}] ${phase}`);
+  });
+
+  try {
+    await runtime.connect();
+    const register = await runtime.registerNode({
+      nodeId,
+      role,
+      capabilities,
+    });
+
+    await publishHeartbeat(runtime, platform);
+
+    console.log(`Connected companion node ${register.node.id} (${platform}, role=${role})`);
+    console.log(`Gateway: ${gatewayUrl}`);
+    console.log(`Capabilities: ${capabilities.join(', ') || '(none)'}`);
+
+    if (options.once) {
+      cleanup();
+      return;
+    }
+
+    heartbeatTimer = setInterval(() => {
+      void publishHeartbeat(runtime, platform).catch((error: unknown) => {
+        const message = error instanceof Error ? error.message : String(error);
+        console.error(`Heartbeat failed: ${message}`);
+      });
+    }, heartbeatSeconds * 1000);
+
+    await new Promise<void>(() => {
+      // Keep process alive until interrupted.
+    });
+  } catch (error) {
+    cleanup();
+    throw error;
+  }
+}
+
+export function registerCompanionCommand(program: Command): void {
+  program
+    .command('companion')
+    .description('Run a minimal gateway companion node client')
+    .option('-c, --config <path>', 'Config file path')
+    .option('--url <url>', 'Gateway WebSocket URL (default from config server.port)')
+    .option('--token <token>', 'Gateway auth token (default from config server.token)')
+    .option('--node-id <id>', 'Node ID to register (default: generated)')
+    .option('--role <role>', 'Node role', 'companion')
+    .option('--platform <platform>', 'Node platform (macos|ios|android|linux|windows|unknown)', 'macos')
+    .option('--capability <name...>', 'Capability list override')
+    .option('--heartbeat <seconds>', 'Heartbeat interval in seconds', '30')
+    .option('--once', 'Connect, register, publish one heartbeat, then exit', false)
+    .action(async (opts: CompanionCommandOptions) => {
+      try {
+        await runCompanionSession(opts);
+      } catch (error) {
+        const message = error instanceof Error ? error.message : String(error);
+        console.error(`Companion command failed: ${message}`);
+        process.exitCode = 1;
+      }
+    });
+}
+
@@ -16,6 +16,7 @@ describe('CLI program', () => {
    expect(commandNames).toContain('backup');
    expect(commandNames).toContain('setup');
    expect(commandNames).toContain('onboard');
+    expect(commandNames).toContain('companion');

    expect(commandNames).toContain('openai-auth');
    expect(commandNames).toContain('openai-key');
@@ -29,6 +29,7 @@ import { registerZaiAuthCommand } from './zai-auth.js';
 import { registerAnthropicAuthCommand } from './anthropic-auth.js';
 import { registerSkillsCommand } from './skills.js';
 import { registerBackupCommand } from './backup.js';
+import { registerCompanionCommand } from './companion.js';

 export function createProgram(): Command {
  const program = new Command();
@@ -58,6 +59,7 @@ export function createProgram(): Command {
  registerAnthropicAuthCommand(program);
  registerSkillsCommand(program);
  registerBackupCommand(program);
+  registerCompanionCommand(program);

  return program;
 }