import { describe, it, expect, beforeAll, afterAll, vi } from 'vitest'; import { resolve } from 'path'; import { createServer } from 'net'; import type { GatewayServerConfig } from '../gateway/server.js'; import { GatewayServer } from '../gateway/server.js'; import { CompanionRuntimeClient, GatewayRpcError } from './runtimeClient.js'; async function canListenOnLocalhost(): Promise { return new Promise((resolvePromise) => { const s = createServer(); s.once('error', () => resolvePromise(false)); s.listen(0, '127.0.0.1', () => { s.close(() => resolvePromise(true)); }); }); } const mockSession = { id: 'test', addMessage: vi.fn(), getHistory: vi.fn(() => []), clear: vi.fn(), setHistory: vi.fn(), replaceHistory: vi.fn(), }; const mockSessionManager = { getSession: vi.fn(() => mockSession), listSessions: vi.fn(() => ['ws:test']), transferSession: vi.fn(), closeSession: vi.fn(), }; const mockModelClient = { chat: vi.fn(async () => ({ content: 'Hello from Flynn!', stopReason: 'end_turn', usage: { inputTokens: 10, outputTokens: 5 }, })), }; const mockToolRegistry = { register: vi.fn(), get: vi.fn(), list: vi.fn(() => []), filteredList: vi.fn(() => []), toAnthropicFormat: vi.fn(() => []), toOpenAIFormat: vi.fn(() => []), filteredToAnthropicFormat: vi.fn(() => []), filteredToOpenAIFormat: vi.fn(() => []), }; const mockToolExecutor = { execute: vi.fn(async () => ({ success: true, output: 'ok' })), }; const TEST_PORT = 18911; const TEST_TOKEN = 'runtime-client-token'; let LISTEN_ALLOWED = true; let server: GatewayServer; beforeAll(async () => { LISTEN_ALLOWED = await canListenOnLocalhost(); if (!LISTEN_ALLOWED) { return; } server = new GatewayServer({ port: TEST_PORT, sessionManager: mockSessionManager as unknown as GatewayServerConfig['sessionManager'], modelClient: mockModelClient, systemPrompt: 'Test prompt', toolRegistry: mockToolRegistry as unknown as GatewayServerConfig['toolRegistry'], toolExecutor: mockToolExecutor as unknown as GatewayServerConfig['toolExecutor'], version: '0.1.0-test', uiDir: resolve(import.meta.dirname, '../gateway/ui'), auth: { token: TEST_TOKEN }, nodes: { enabled: true, allowedRoles: ['companion'], featureGates: { 'ui.canvas': true }, locationEnabled: true, pushEnabled: true, }, }); await server.start(); }); afterAll(async () => { if (!LISTEN_ALLOWED) { return; } await server.stop(); }); describe('CompanionRuntimeClient', () => { it('validates requestTimeoutMs option', () => { expect(() => { new CompanionRuntimeClient({ url: 'ws://127.0.0.1:1', requestTimeoutMs: 0, }); }).toThrow('requestTimeoutMs must be a positive number'); }); it('dispatches gateway events to subscribed handlers and supports unsubscribe', () => { const client = new CompanionRuntimeClient({ url: 'ws://127.0.0.1:1', }); const handler = vi.fn(); const unsubscribe = client.subscribeEvents(handler); (client as unknown as { handleMessage: (raw: string) => void }).handleMessage( JSON.stringify({ id: 42, event: 'agent.stream', data: { token: 'hello' }, }), ); expect(handler).toHaveBeenCalledWith('agent.stream', { token: 'hello' }); unsubscribe(); (client as unknown as { handleMessage: (raw: string) => void }).handleMessage( JSON.stringify({ id: 43, event: 'agent.stream', data: { token: 'world' }, }), ); expect(handler).toHaveBeenCalledTimes(1); }); it('isolates subscriber callback failures', () => { const client = new CompanionRuntimeClient({ url: 'ws://127.0.0.1:1', }); const badHandler = vi.fn(() => { throw new Error('subscriber failed'); }); const goodHandler = vi.fn(); client.subscribeEvents(badHandler); client.subscribeEvents(goodHandler); expect(() => { (client as unknown as { handleMessage: (raw: string) => void }).handleMessage( JSON.stringify({ id: 44, event: 'agent.stream', data: { token: 'safe' }, }), ); }).not.toThrow(); expect(badHandler).toHaveBeenCalledOnce(); expect(goodHandler).toHaveBeenCalledWith('agent.stream', { token: 'safe' }); }); it('supports filtered subscribeEvent helper', () => { const client = new CompanionRuntimeClient({ url: 'ws://127.0.0.1:1', }); const streamHandler = vi.fn(); const unsubscribe = client.subscribeEvent<{ token: string }>('agent.stream', streamHandler); (client as unknown as { handleMessage: (raw: string) => void }).handleMessage( JSON.stringify({ id: 45, event: 'agent.stream', data: { token: 'first' }, }), ); (client as unknown as { handleMessage: (raw: string) => void }).handleMessage( JSON.stringify({ id: 46, event: 'agent.typing', data: { active: true }, }), ); expect(streamHandler).toHaveBeenCalledTimes(1); expect(streamHandler).toHaveBeenCalledWith({ token: 'first' }); unsubscribe(); (client as unknown as { handleMessage: (raw: string) => void }).handleMessage( JSON.stringify({ id: 47, event: 'agent.stream', data: { token: 'second' }, }), ); expect(streamHandler).toHaveBeenCalledTimes(1); }); it('supports subscribeAgentStream and subscribeAgentTyping helpers', () => { const client = new CompanionRuntimeClient({ url: 'ws://127.0.0.1:1', }); const streamHandler = vi.fn(); const typingHandler = vi.fn(); client.subscribeAgentStream(streamHandler); client.subscribeAgentTyping(typingHandler); (client as unknown as { handleMessage: (raw: string) => void }).handleMessage( JSON.stringify({ id: 52, event: 'agent.stream', data: { token: 'x' }, }), ); (client as unknown as { handleMessage: (raw: string) => void }).handleMessage( JSON.stringify({ id: 53, event: 'agent.typing', data: { active: true }, }), ); expect(streamHandler).toHaveBeenCalledWith({ token: 'x' }); expect(typingHandler).toHaveBeenCalledWith({ active: true }); }); it('supports subscribeContextWarning helper', () => { const client = new CompanionRuntimeClient({ url: 'ws://127.0.0.1:1', }); const warningHandler = vi.fn(); client.subscribeContextWarning(warningHandler); (client as unknown as { handleMessage: (raw: string) => void }).handleMessage( JSON.stringify({ id: 56, event: 'context_warning', data: { thresholdPct: 80, estimatedPct: 92 }, }), ); expect(warningHandler).toHaveBeenCalledWith({ thresholdPct: 80, estimatedPct: 92 }); }); it('clears all event subscriptions', () => { const client = new CompanionRuntimeClient({ url: 'ws://127.0.0.1:1', }); const handlerA = vi.fn(); const handlerB = vi.fn(); client.subscribeEvents(handlerA); client.subscribeEvent('agent.stream', handlerB); client.clearEventSubscriptions(); (client as unknown as { handleMessage: (raw: string) => void }).handleMessage( JSON.stringify({ id: 50, event: 'agent.stream', data: { token: 'cleared' }, }), ); expect(handlerA).not.toHaveBeenCalled(); expect(handlerB).not.toHaveBeenCalled(); }); it('dispose clears subscriptions and is safe to call repeatedly', () => { const client = new CompanionRuntimeClient({ url: 'ws://127.0.0.1:1', }); const handler = vi.fn(); client.subscribeEvents(handler); client.dispose(); client.dispose(); (client as unknown as { handleMessage: (raw: string) => void }).handleMessage( JSON.stringify({ id: 51, event: 'agent.stream', data: { token: 'after-dispose' }, }), ); expect(handler).not.toHaveBeenCalled(); }); it('waitForEvent resolves using optional predicate filter', async () => { const client = new CompanionRuntimeClient({ url: 'ws://127.0.0.1:1', }); const awaited = client.waitForEvent<{ seq: number }>('agent.stream', { timeoutMs: 2000, predicate: (data) => data.seq === 2, }); (client as unknown as { handleMessage: (raw: string) => void }).handleMessage( JSON.stringify({ id: 48, event: 'agent.stream', data: { seq: 1 }, }), ); (client as unknown as { handleMessage: (raw: string) => void }).handleMessage( JSON.stringify({ id: 49, event: 'agent.stream', data: { seq: 2 }, }), ); await expect(awaited).resolves.toEqual({ seq: 2 }); }); it('waitForEvent rejects on timeout', async () => { vi.useFakeTimers(); const client = new CompanionRuntimeClient({ url: 'ws://127.0.0.1:1', }); const awaited = expect( client.waitForEvent('agent.stream', { timeoutMs: 100 }), ).rejects.toThrow('Timed out waiting for event agent.stream'); await vi.advanceTimersByTimeAsync(100); await awaited; vi.useRealTimers(); }); it('waitForEvent supports AbortSignal cancellation', async () => { const client = new CompanionRuntimeClient({ url: 'ws://127.0.0.1:1', }); const controller = new AbortController(); const awaited = expect( client.waitForEvent('agent.stream', { signal: controller.signal, timeoutMs: 10_000 }), ).rejects.toThrow('Aborted while waiting for event agent.stream'); controller.abort(); await awaited; }); it('waitForAgentStream resolves on agent.stream events', async () => { const client = new CompanionRuntimeClient({ url: 'ws://127.0.0.1:1', }); const awaited = client.waitForAgentStream<{ token: string }>({ timeoutMs: 2000 }); (client as unknown as { handleMessage: (raw: string) => void }).handleMessage( JSON.stringify({ id: 54, event: 'agent.stream', data: { token: 'typed-stream' }, }), ); await expect(awaited).resolves.toEqual({ token: 'typed-stream' }); }); it('waitForAgentTyping resolves on agent.typing events', async () => { const client = new CompanionRuntimeClient({ url: 'ws://127.0.0.1:1', }); const awaited = client.waitForAgentTyping<{ active: boolean }>({ timeoutMs: 2000 }); (client as unknown as { handleMessage: (raw: string) => void }).handleMessage( JSON.stringify({ id: 55, event: 'agent.typing', data: { active: true }, }), ); await expect(awaited).resolves.toEqual({ active: true }); }); it('waitForContextWarning resolves on context_warning events', async () => { const client = new CompanionRuntimeClient({ url: 'ws://127.0.0.1:1', }); const awaited = client.waitForContextWarning<{ thresholdPct: number; estimatedPct: number }>({ timeoutMs: 2000, }); (client as unknown as { handleMessage: (raw: string) => void }).handleMessage( JSON.stringify({ id: 57, event: 'context_warning', data: { thresholdPct: 75, estimatedPct: 88 }, }), ); await expect(awaited).resolves.toEqual({ thresholdPct: 75, estimatedPct: 88 }); }); it('connects and performs node registration + capability discovery', async () => { if (!LISTEN_ALLOWED) { return; } const client = new CompanionRuntimeClient({ url: `ws://127.0.0.1:${TEST_PORT}`, token: TEST_TOKEN, }); await client.connect(); expect(client.connected).toBe(true); try { const register = await client.registerNode({ nodeId: 'macos-companion', role: 'companion', capabilities: ['ui.canvas', 'node.location.read'], }); expect(register.registered).toBe(true); expect(register.node.id).toBe('macos-companion'); expect(register.protocol.serverVersion).toBe(1); expect(register.capabilities.enabled).toContain('ui.canvas'); const nodeCaps = await client.getNodeCapabilities(); expect(nodeCaps.node.id).toBe('macos-companion'); expect(nodeCaps.capabilities.enabled).toContain('ui.canvas'); const systemCaps = await client.getSystemCapabilities(); expect(systemCaps.nodes.enabled).toBe(true); expect(systemCaps.nodes.registered).toBe(true); expect(systemCaps.nodes.nodeId).toBe('macos-companion'); } finally { client.disconnect(); } }); it('bootstraps node registration and capabilities in one call', async () => { if (!LISTEN_ALLOWED) { return; } const client = new CompanionRuntimeClient({ url: `ws://127.0.0.1:${TEST_PORT}`, token: TEST_TOKEN, }); await client.connect(); try { const boot = await client.bootstrapNode( { nodeId: 'bootstrap-runtime-node', role: 'companion', capabilities: ['ui.canvas'], }, { includeSystemCapabilities: true }, ); expect(boot.register.registered).toBe(true); expect(boot.capabilities.node.id).toBe('bootstrap-runtime-node'); expect(boot.systemCapabilities?.nodes.enabled).toBe(true); expect(boot.systemCapabilities?.nodes.registered).toBe(true); } finally { client.disconnect(); } }); it('updates status/location/push and exposes them through system.nodes', async () => { if (!LISTEN_ALLOWED) { return; } const client = new CompanionRuntimeClient({ url: `ws://127.0.0.1:${TEST_PORT}`, token: TEST_TOKEN, }); await client.connect(); try { await client.registerNode({ nodeId: 'ios-companion', role: 'companion', capabilities: ['node.location.write', 'node.push.register'], }); const status = await client.setNodeStatus({ platform: 'ios', appVersion: '1.2.3', batteryPct: 77, powerSource: 'battery', }); expect(status.updated).toBe(true); expect(status.status.platform).toBe('ios'); const location = await client.setNodeLocation({ latitude: 37.78, longitude: -122.41, source: 'gps', }); expect(location.updated).toBe(true); expect(location.location.source).toBe('gps'); const locationGet = await client.getNodeLocation(); expect(locationGet.location?.latitude).toBe(37.78); const push = await client.setNodePushToken({ provider: 'apns', token: 'abcdef1234567890abcdef1234567890abcdef1234567890abcdef1234567890', topic: 'dev.flynn.companion', environment: 'production', }); expect(push.updated).toBe(true); expect(push.push.provider).toBe('apns'); expect(push.push.tokenPreview.startsWith('***')).toBe(true); const nodes = await client.listSystemNodes({ role: 'companion', platform: 'ios' }); expect(nodes.summary.total).toBeGreaterThanOrEqual(1); const current = nodes.nodes.find((entry) => entry.nodeId === 'ios-companion'); expect(current).toBeTruthy(); expect(current?.status?.platform).toBe('ios'); expect(current?.push?.provider).toBe('apns'); } finally { client.disconnect(); } }); it('surfaces gateway errors as GatewayRpcError', async () => { if (!LISTEN_ALLOWED) { return; } const client = new CompanionRuntimeClient({ url: `ws://127.0.0.1:${TEST_PORT}`, token: TEST_TOKEN, }); await client.connect(); try { await expect(client.registerNode({ nodeId: 'observer-node', role: 'observer', capabilities: ['node.location.read'], })).rejects.toBeInstanceOf(GatewayRpcError); await expect(client.registerNode({ nodeId: 'observer-node', role: 'observer', capabilities: ['node.location.read'], })).rejects.toMatchObject({ code: -5, }); } finally { client.disconnect(); } }); it('supports canvas artifact lifecycle via typed helper methods', async () => { if (!LISTEN_ALLOWED) { return; } const client = new CompanionRuntimeClient({ url: `ws://127.0.0.1:${TEST_PORT}`, token: TEST_TOKEN, }); await client.connect(); try { const sessionId = 'ws:companion-canvas'; const put = await client.putCanvasArtifact({ sessionId, artifactId: 'artifact-1', type: 'markdown', title: 'Companion note', content: { body: '# Hello' }, metadata: { source: 'runtime-client-test' }, }); expect(put.upserted).toBe(true); expect(put.artifact.id).toBe('artifact-1'); expect(put.artifact.type).toBe('markdown'); const list = await client.listCanvasArtifacts(sessionId); expect(list.artifacts.length).toBeGreaterThanOrEqual(1); expect(list.artifacts.some((artifact) => artifact.id === 'artifact-1')).toBe(true); const get = await client.getCanvasArtifact({ sessionId, artifactId: 'artifact-1' }); expect(get.artifact.id).toBe('artifact-1'); expect(get.artifact.title).toBe('Companion note'); const del = await client.deleteCanvasArtifact({ sessionId, artifactId: 'artifact-1' }); expect(del.deleted).toBe(true); const clear = await client.clearCanvasArtifacts(sessionId); expect(clear.cleared).toBe(0); } finally { client.disconnect(); } }); it('supports autoConnect mode for one-shot RPC usage', async () => { if (!LISTEN_ALLOWED) { return; } const client = new CompanionRuntimeClient({ url: `ws://127.0.0.1:${TEST_PORT}`, token: TEST_TOKEN, autoConnect: true, }); expect(client.connected).toBe(false); try { const register = await client.registerNode({ nodeId: 'auto-connect-node', role: 'companion', capabilities: ['ui.canvas'], }); expect(register.registered).toBe(true); expect(client.connected).toBe(true); } finally { client.disconnect(); } }); });