vdm/extension/tests/transport/websocket.test.ts

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<WebSocketTransportDeps> = {}) {
  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('the pairing token survives a browser restart: a fresh transport instance over the same storage reuses it, no re-pairing', async () => {
    const port = nextPort();
    daemon = await FakeDaemon.start({ port, acceptToken: null });
    const h = memDeps(); // stands in for browser.storage.local, which outlives the page
    transport = makeTransport(port, h);

    await transport.connect();
    expect(daemon.pairCount).toBe(1);
    expect(h.store.token).toBe('tok-issued-1');
    transport.disconnect();

    // Simulate "the browser restarted": a brand new transport instance, same backing
    // store (in reality, the same on-disk profile), no in-memory state carried over.
    const restarted = makeTransport(port, h);
    await restarted.connect();
    try {
      expect(restarted.state).toBe('connected');
      expect(daemon.pairCount).toBe(1); // still just the one pairing, ever
      const lastHello = [...daemon.seen].reverse().find((s) => s.method === 'session.hello');
      expect(lastHello?.token).toBe('tok-issued-1');
    } finally {
      restarted.disconnect();
    }
  });

  it('pairWithCode drops any stale token and pairs fresh with the typed code', async () => {
    const port = nextPort();
    daemon = await FakeDaemon.start({ port, acceptToken: 'stale' });
    const h = memDeps({ token: 'stale' });
    transport = makeTransport(port, h);

    await transport.pairWithCode('4821');

    expect(transport.state).toBe('connected');
    expect(daemon.pairCount).toBe(1);
    const pairCall = daemon.seen.find((s) => s.method === 'session.pair');
    expect((pairCall?.params as { code?: string }).code).toBe('4821');
    expect(h.store.token).toBe('tok-issued-1');
  });

  it('unpair clears the stored token and disconnects', 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();

    await transport.unpair();

    expect(h.store.token).toBeNull();
    expect(transport.state).toBe('disconnected');
  });

  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');
  });
});