import { afterEach, describe, expect, it } from 'vitest'; import { WebSocket as WsClient } from 'ws'; import { MethodNotAllowedError, RpcError, RpcTimeoutError, TransportClosedError, } from '../../src/background/transport/types.js'; import { WebSocketTransport, type WebSocketCtor, type WebSocketTransportDeps, } from '../../src/background/transport/websocket.js'; import { FakeDaemon, nextPort, tick } from './helpers.js'; const CTOR = WsClient as unknown as WebSocketCtor; interface DepsHandle { deps: WebSocketTransportDeps; store: { token: string | null; port: number | null }; } function memDeps(init: Partial<{ token: string | null }> = {}): DepsHandle { const store = { token: init.token ?? null, port: null as number | null }; return { store, deps: { getToken: async () => store.token, setToken: async (t) => { store.token = t; }, getCachedPort: async () => store.port, setCachedPort: async (p) => { store.port = p; }, extensionId: '11111111-2222-3333-4444-555555555555', }, }; } function makeTransport(port: number, h: DepsHandle, extra: Partial = {}) { return new WebSocketTransport({ ...h.deps, ...extra, webSocketCtor: CTOR, portRange: { start: port, end: port }, openTimeoutMs: 500, backoff: { baseMs: 20, factor: 2, maxMs: 120, jitter: 0 }, }); } let daemon: FakeDaemon | undefined; let transport: WebSocketTransport | undefined; afterEach(async () => { transport?.disconnect(); transport = undefined; await daemon?.stop(); daemon = undefined; }); describe('WebSocketTransport', () => { it('discovers the port, completes session.hello, and reports daemon info', async () => { const port = nextPort(); daemon = await FakeDaemon.start({ port, acceptToken: 'good-token' }); const h = memDeps({ token: 'good-token' }); transport = makeTransport(port, h); await transport.connect(); expect(transport.state).toBe('connected'); expect(transport.status.daemonVersion).toBe('1.2.3-fake'); expect(transport.status.capabilities).toEqual(['media', 'grabber']); expect(transport.status.sessionId).toBe('sess-1'); expect(h.store.port).toBe(port); // cached for next time expect(daemon.pairCount).toBe(0); // had a valid token, no pairing needed }); it('pairs when unpaired, stores the token, and re-hellos with it', async () => { const port = nextPort(); daemon = await FakeDaemon.start({ port, acceptToken: null }); const h = memDeps(); transport = makeTransport(port, h); await transport.connect(); expect(transport.state).toBe('connected'); expect(daemon.pairCount).toBe(1); expect(h.store.token).toBe('tok-issued-1'); // persisted expect(daemon.helloCount).toBe(2); // once unpaired, once with the fresh token const lastHello = [...daemon.seen].reverse().find((s) => s.method === 'session.hello'); expect(lastHello?.token).toBe('tok-issued-1'); }); it('with autoPair off, a wrong token surfaces needsPairing and does NOT retry', async () => { const port = nextPort(); daemon = await FakeDaemon.start({ port, acceptToken: 'the-real-one' }); const h = memDeps({ token: 'stale-token' }); transport = makeTransport(port, h, { autoPair: false }); await expect(transport.connect()).rejects.toBeInstanceOf(RpcError); expect(transport.state).toBe('disconnected'); expect(transport.status.needsPairing).toBe(true); const helloAfterReject = daemon.helloCount; await tick(150); // longer than the (tiny) backoff would be expect(daemon.helloCount).toBe(helloAfterReject); // no reconnect storm }); it('treats a pairing rate-limit as needsPairing with a retry hint, no retry', async () => { const port = nextPort(); daemon = await FakeDaemon.start({ port, acceptToken: null, pairError: { code: -32014, message: 'locked out', data: { retryAfterSec: 60 } }, }); transport = makeTransport(port, memDeps()); await expect(transport.connect()).rejects.toBeInstanceOf(RpcError); expect(transport.status.needsPairing).toBe(true); expect(transport.status.retryAfterSec).toBe(60); const n = daemon.helloCount; await tick(150); expect(daemon.helloCount).toBe(n); }); it('treats a protocol-major mismatch as fatal and stops retrying', async () => { const port = nextPort(); daemon = await FakeDaemon.start({ port, versionMismatch: true }); transport = makeTransport(port, memDeps({ token: 'x' })); await expect(transport.connect()).rejects.toBeInstanceOf(RpcError); expect(transport.status.fatal).toMatch(/protocol mismatch/); const n = daemon.helloCount; await tick(150); expect(daemon.helloCount).toBe(n); }); it('forwards a call and returns the result', async () => { const port = nextPort(); daemon = await FakeDaemon.start({ port, acceptToken: 't' }); transport = makeTransport(port, memDeps({ token: 't' })); await transport.connect(); await expect(transport.call('download.list', { filter: null })).resolves.toEqual({ total: 0, items: [], }); }); it('rejects a call with RpcTimeoutError when the daemon never answers', async () => { const port = nextPort(); daemon = await FakeDaemon.start({ port, acceptToken: 't', blackhole: ['download.list'] }); transport = makeTransport(port, memDeps({ token: 't' })); await transport.connect(); await expect( transport.call('download.list', { filter: null }, { timeoutMs: 40 }), ).rejects.toBeInstanceOf(RpcTimeoutError); }); it('rejects a uds-only method locally, before it hits the wire', async () => { const port = nextPort(); daemon = await FakeDaemon.start({ port, acceptToken: 't' }); transport = makeTransport(port, memDeps({ token: 't' })); await transport.connect(); const before = daemon.seen.length; await expect( transport.call('settings.set', { values: { 'capture.enabled': false } }), ).rejects.toBeInstanceOf(MethodNotAllowedError); expect(daemon.seen.length).toBe(before); }); it('delivers server notifications to on() listeners', async () => { const port = nextPort(); daemon = await FakeDaemon.start({ port, acceptToken: 't' }); transport = makeTransport(port, memDeps({ token: 't' })); await transport.connect(); const seen: unknown[] = []; transport.on('event.task.progress', (p) => seen.push(p)); daemon.notify('event.task.progress', { tasks: [{ taskId: 'a', downloadedBytes: 1 }] }); await tick(20); expect(seen).toEqual([{ tasks: [{ taskId: 'a', downloadedBytes: 1 }] }]); }); it('fails in-flight calls on a drop and then reconnects on its own', async () => { const port = nextPort(); daemon = await FakeDaemon.start({ port, acceptToken: 't', blackhole: ['download.list'] }); transport = makeTransport(port, memDeps({ token: 't' })); await transport.connect(); const pending = transport.call('download.list', { filter: null }, { timeoutMs: 5000 }); daemon.dropAll(); await expect(pending).rejects.toBeInstanceOf(TransportClosedError); expect(transport.state).toBe('disconnected'); // backoff base is 20ms; give it room to come back for (let i = 0; i < 40 && transport.state !== 'connected'; i += 1) await tick(15); expect(transport.state).toBe('connected'); }); it('disconnect() stops all reconnection', async () => { const port = nextPort(); daemon = await FakeDaemon.start({ port, acceptToken: 't' }); transport = makeTransport(port, memDeps({ token: 't' })); await transport.connect(); transport.disconnect(); daemon.dropAll(); const n = daemon.helloCount; await tick(150); expect(daemon.helloCount).toBe(n); expect(transport.state).toBe('disconnected'); }); it('when nothing answers, connect() rejects but keeps retrying in the background', async () => { const port = nextPort(); // no daemon on it const h = memDeps({ token: 't' }); transport = makeTransport(port, h); await expect(transport.connect()).rejects.toBeInstanceOf(TransportClosedError); expect(transport.state).toBe('disconnected'); expect(transport.status.fatal).toBeNull(); expect(transport.status.needsPairing).toBe(false); // now bring a daemon up on that port; the transport should find it daemon = await FakeDaemon.start({ port, acceptToken: 't' }); for (let i = 0; i < 60 && transport.state !== 'connected'; i += 1) await tick(15); expect(transport.state).toBe('connected'); }); });