Shade/packages/shade-transport-webrtc/tests/connection.test.ts

import { afterEach, describe, expect, it } from 'bun:test';
import {
  encodeChunkFrame,
  encodeResumeQueryFrame,
  randomRequestId,
  STREAM_ID_LEN,
  WIRE_CHUNK,
  WIRE_CHUNK_ACK,
  WIRE_RESUME_QUERY,
  WIRE_RESUME_STATE,
  WIRE_ERROR,
  bytesEqual,
} from '../src/wire.js';
import { MemoryRtcFactory } from '../src/memory-rtc.js';
import { MemoryShadeBridge, WebRtcSignalingChannel } from '../src/signaling.js';
import { WebRtcConnectionManager } from '../src/manager.js';

afterEach(() => {
  MemoryRtcFactory.reset();
});

async function setupPair(opts?: {
  bobReceiver?: import('../src/connection.js').WebRtcReceiverHooks;
  aliceReceiver?: import('../src/connection.js').WebRtcReceiverHooks;
}): Promise<{
  alice: WebRtcConnectionManager;
  bob: WebRtcConnectionManager;
  aliceSig: WebRtcSignalingChannel;
  bobSig: WebRtcSignalingChannel;
}> {
  const { a, b } = MemoryShadeBridge.linked('alice', 'bob');
  const aliceSig = new WebRtcSignalingChannel(a);
  const bobSig = new WebRtcSignalingChannel(b);
  const factory = new MemoryRtcFactory();
  const alice = new WebRtcConnectionManager({
    factory,
    signaling: aliceSig,
    ...(opts?.aliceReceiver !== undefined ? { receiver: opts.aliceReceiver } : {}),
  });
  const bob = new WebRtcConnectionManager({
    factory,
    signaling: bobSig,
    ...(opts?.bobReceiver !== undefined ? { receiver: opts.bobReceiver } : {}),
  });
  return { alice, bob, aliceSig, bobSig };
}

describe('WebRtcConnection — caller/callee handshake', () => {
  it('opens a data channel after offer/answer/ICE flow', async () => {
    const { alice, bob } = await setupPair();
    const conn = await alice.getOrCreate('bob');
    expect(conn.state).toBe('connected');
    expect(bob.isConnected('alice')).toBe(true);

    // The peer connection on the bob side should also be open and reachable.
    const bobConn = await bob.getOrCreate('alice');
    expect(bobConn.state).toBe('connected');

    alice.destroy();
    bob.destroy();
  });

  it('routes a chunk request to the receiver hook and replies with chunk-ack', async () => {
    const calls: Array<{
      from: string;
      streamId: string;
      laneId: number;
      seq: bigint;
      bytes: number;
    }> = [];
    const { alice, bob } = await setupPair({
      bobReceiver: {
        async onChunk(from, streamId, laneId, seq, envelope) {
          calls.push({ from, streamId, laneId, seq, bytes: envelope.length });
          return { lastSeq: Number(seq), bytesReceived: envelope.length };
        },
        async onResumeQuery() {
          return null;
        },
      },
    });

    const conn = await alice.getOrCreate('bob');
    const requestId = randomRequestId();
    const streamId = new Uint8Array(STREAM_ID_LEN).fill(0xaa);
    const envelope = new Uint8Array(64);
    for (let i = 0; i < envelope.length; i++) envelope[i] = i;

    const frame = encodeChunkFrame({
      type: WIRE_CHUNK,
      requestId,
      streamId,
      laneId: 3,
      seq: 7n,
      envelope,
    });
    const response = await conn.request(frame, requestId);
    expect(response.type).toBe(WIRE_CHUNK_ACK);
    if (response.type === WIRE_CHUNK_ACK) {
      expect(bytesEqual(response.requestId, requestId)).toBe(true);
      expect(response.lastSeq).toBe(7);
      expect(response.bytesReceived).toBe(64);
    }
    expect(calls).toHaveLength(1);
    expect(calls[0]!.laneId).toBe(3);
    expect(calls[0]!.seq).toBe(7n);
    expect(calls[0]!.bytes).toBe(64);

    alice.destroy();
    bob.destroy();
  });

  it('returns error frame when receiver throws', async () => {
    const { alice, bob } = await setupPair({
      bobReceiver: {
        async onChunk() {
          throw new Error('nope');
        },
        async onResumeQuery() {
          return null;
        },
      },
    });

    const conn = await alice.getOrCreate('bob');
    const requestId = randomRequestId();
    const streamId = new Uint8Array(STREAM_ID_LEN).fill(0x01);
    const frame = encodeChunkFrame({
      type: WIRE_CHUNK,
      requestId,
      streamId,
      laneId: 0,
      seq: 0n,
      envelope: new Uint8Array(4),
    });
    const response = await conn.request(frame, requestId);
    expect(response.type).toBe(WIRE_ERROR);
    if (response.type === WIRE_ERROR) {
      expect(response.json).toContain('nope');
    }

    alice.destroy();
    bob.destroy();
  });

  it('handles resume-query with not-found → error frame', async () => {
    const { alice, bob } = await setupPair({
      bobReceiver: {
        async onChunk() {
          return { lastSeq: 0 };
        },
        async onResumeQuery() {
          return null;
        },
      },
    });
    const conn = await alice.getOrCreate('bob');
    const requestId = randomRequestId();
    const streamId = new Uint8Array(STREAM_ID_LEN).fill(0x55);
    const frame = encodeResumeQueryFrame({
      type: WIRE_RESUME_QUERY,
      requestId,
      streamId,
    });
    const res = await conn.request(frame, requestId);
    expect(res.type).toBe(WIRE_ERROR);
    if (res.type === WIRE_ERROR) {
      expect(res.json).toContain('not found');
    }
    alice.destroy();
    bob.destroy();
  });

  it('handles resume-query that returns state → resume-state frame', async () => {
    const { alice, bob } = await setupPair({
      bobReceiver: {
        async onChunk() {
          return { lastSeq: 0 };
        },
        async onResumeQuery(_from, streamId) {
          return { streamId, lanes: [{ laneId: 0, lastSeqAcked: 11 }] };
        },
      },
    });
    const conn = await alice.getOrCreate('bob');
    const requestId = randomRequestId();
    const streamId = new Uint8Array(STREAM_ID_LEN).fill(0x77);
    const frame = encodeResumeQueryFrame({
      type: WIRE_RESUME_QUERY,
      requestId,
      streamId,
    });
    const res = await conn.request(frame, requestId);
    expect(res.type).toBe(WIRE_RESUME_STATE);
    if (res.type === WIRE_RESUME_STATE) {
      const parsed = JSON.parse(res.json) as { lanes: { laneId: number; lastSeqAcked: number }[] };
      expect(parsed.lanes[0]!.lastSeqAcked).toBe(11);
    }
    alice.destroy();
    bob.destroy();
  });
});

describe('WebRtcConnectionManager pool', () => {
  it('reuses one connection per peer', async () => {
    const { alice, bob } = await setupPair();
    const c1 = await alice.getOrCreate('bob');
    const c2 = await alice.getOrCreate('bob');
    expect(c1).toBe(c2);
    alice.destroy();
    bob.destroy();
  });

  it('removes the connection from the pool when it closes', async () => {
    const { alice, bob } = await setupPair();
    const conn = await alice.getOrCreate('bob');
    expect(alice.isConnected('bob')).toBe(true);
    await conn.close('test');
    expect(alice.isConnected('bob')).toBe(false);
    alice.destroy();
    bob.destroy();
  });
});