web-terminal/term-relay/test/mux-session.test.ts

import { describe, test, expect, vi } from 'vitest'
import { createMuxSession, type MuxSession, type MuxStreamHandle } from '../mux/mux-session.js'
import { encodeMuxFrame, decodeHeader, type MuxOpen } from '../mux/frame-codec.js'
import type { ScheduleHandle } from '../mux/heartbeat.js'

const noAutoSchedule = (): { schedule: (fn: () => void) => ScheduleHandle; tick: () => void } => {
  let scheduled: (() => void) | null = null
  return {
    schedule: (fn) => {
      scheduled = fn
      return { cancel: () => (scheduled = null) }
    },
    tick: () => scheduled?.(),
  }
}

function openFor(streamId: number): MuxOpen {
  return {
    streamId,
    subdomain: 'alice',
    requestPath: '/term',
    originHeader: 'https://alice.term.example.com',
    remoteAddrHash: 'hash',
    capabilityTokenRef: 'jti-1',
  }
}

interface Pair {
  relay: MuxSession
  agent: MuxSession
  opened: { open: MuxOpen; stream: MuxStreamHandle }[]
  agentData: { streamId: number; payload: Uint8Array }[]
}

function wirePair(relayWindow = 1_000_000, agentWindow = 1_000_000): Pair {
  const opened: Pair['opened'] = []
  const agentData: Pair['agentData'] = []
  let agent!: MuxSession
  const relay = createMuxSession({
    role: 'relay',
    sendWire: (f) => agent.onWire(f),
    onOpen: () => {},
    onData: () => {},
    onClose: () => {},
    onDead: () => {},
    maxFrameBytes: 1024,
    initialWindowBytes: relayWindow,
    schedule: noAutoSchedule().schedule,
  })
  agent = createMuxSession({
    role: 'agent',
    sendWire: (f) => relay.onWire(f),
    onOpen: (open, stream) => opened.push({ open, stream }),
    onData: (streamId, payload) => agentData.push({ streamId, payload }),
    onClose: () => {},
    onDead: () => {},
    maxFrameBytes: 1024,
    initialWindowBytes: agentWindow,
    schedule: noAutoSchedule().schedule,
  })
  return { relay, agent, opened, agentData }
}

describe('mux session (T6)', () => {
  test('openStream → agent onOpen with a fresh monotonic streamId; never reused', () => {
    const p = wirePair()
    p.relay.openStream(openFor(0))
    p.relay.openStream(openFor(0))
    expect(p.opened.map((o) => o.open.streamId)).toEqual([1, 2])
  })

  test('writeData delivers a byte-identical OPAQUE payload (INV2), never inspected', () => {
    const p = wirePair()
    const h = p.relay.openStream(openFor(0))
    const marker = new TextEncoder().encode('PLAINTEXT-MARKER-🔒')
    expect(h.writeData(marker)).toBe(true)
    expect(p.agentData).toHaveLength(1)
    expect([...p.agentData[0]!.payload]).toEqual([...marker])
  })

  test('backpressure: writeData returns false when the window is full; WINDOW_UPDATE resumes', () => {
    // relay send window = 10, agent receive window huge (no auto-replenish crossing threshold).
    const p = wirePair(10, 1_000_000)
    const h = p.relay.openStream(openFor(0))
    const streamId = h.streamId
    const a = new Uint8Array(8).fill(1)
    const b = new Uint8Array(8).fill(2)
    expect(h.writeData(a)).toBe(true) // window 10 → 2
    expect(h.writeData(b)).toBe(false) // 8 > 2 → buffered
    expect(p.agentData).toHaveLength(1)
    // Simulate a peer WINDOW_UPDATE granting credit → buffered `b` flushes.
    p.relay.onWire(encodeWU(streamId, 50))
    expect(p.agentData).toHaveLength(2)
    expect([...p.agentData[1]!.payload]).toEqual([...b])
  })

  test('per-stream isolation: exhausting stream A does not block stream B', () => {
    const p = wirePair(10, 1_000_000)
    const a = p.relay.openStream(openFor(0))
    const b = p.relay.openStream(openFor(0))
    expect(a.writeData(new Uint8Array(10).fill(1))).toBe(true) // A window exhausted
    expect(a.writeData(new Uint8Array(1))).toBe(false)
    expect(b.writeData(new Uint8Array(10).fill(2))).toBe(true) // B unaffected
  })

  test('inbound DATA for an unknown stream → CLOSE+RST for that stream only (tunnel stays up)', () => {
    const wire: Uint8Array[] = []
    const s = createMuxSession({
      role: 'agent',
      sendWire: (f) => wire.push(f),
      onOpen: () => {},
      onData: () => {},
      onClose: () => {},
      onDead: () => {},
      maxFrameBytes: 1024,
      initialWindowBytes: 1000,
      schedule: noAutoSchedule().schedule,
    })
    s.onWire(encodeMuxFrame({ version: 1, type: 'data', fin: false, rst: false, streamId: 99, payloadLen: 1 }, new Uint8Array([7])))
    const rst = wire.map(decodeHeader).find((x) => x.type === 'close' && x.rst)
    expect(rst?.streamId).toBe(99)
  })

  test('frame ceiling: inbound payloadLen > maxFrameBytes → RST that stream, no OOM', () => {
    const wire: Uint8Array[] = []
    const s = createMuxSession({
      role: 'agent',
      sendWire: (f) => wire.push(f),
      onOpen: () => {},
      onData: () => {},
      onClose: () => {},
      onDead: () => {},
      maxFrameBytes: 4,
      initialWindowBytes: 1000,
      schedule: noAutoSchedule().schedule,
    })
    const oversized = encodeMuxFrame({ version: 1, type: 'data', fin: false, rst: false, streamId: 5, payloadLen: 5 }, new Uint8Array(5))
    s.onWire(oversized)
    const rst = wire.map(decodeHeader).find((x) => x.rst)
    expect(rst?.streamId).toBe(5)
  })

  test('heartbeat: no PONG for the miss limit ⇒ onDead', () => {
    const sched = noAutoSchedule()
    const onDead = vi.fn()
    const s = createMuxSession({
      role: 'relay',
      sendWire: () => {}, // peer never responds
      onOpen: () => {},
      onData: () => {},
      onClose: () => {},
      onDead,
      maxFrameBytes: 1024,
      initialWindowBytes: 1000,
      schedule: sched.schedule,
    })
    s.openStream(openFor(0)) // starts heartbeat (ping #1)
    sched.tick() // miss 1
    sched.tick() // miss 2
    sched.tick() // miss 3 → dead
    expect(onDead).toHaveBeenCalledTimes(1)
  })
})

function encodeWU(streamId: number, credit: number): Uint8Array {
  const payload = new Uint8Array(4)
  new DataView(payload.buffer).setUint32(0, credit, false)
  return encodeMuxFrame(
    { version: 1, type: 'windowUpdate', fin: false, rst: false, streamId, payloadLen: 4 },
    payload,
  )
}