agent-desk-ng/web/lib/realtime-connection.test.mjs

// Modified for hardening and realtime recovery; see docs/hardening/README.md and docs/fixes/chat-realtime-recovery.md.
import assert from "node:assert/strict"
import test from "node:test"

import { createRealtimeConnectionManager } from "./realtime-connection.ts"

// Minimal WebSocket mock for testing connection manager lifecycle
class MockWebSocket {
  static CONNECTING = 0
  static OPEN = 1
  static CLOSING = 2
  static CLOSED = 3

  constructor() {
    this.readyState = MockWebSocket.CONNECTING
    this.listeners = new Map()
  }

  addEventListener(type, listener) {
    if (!this.listeners.has(type)) {
      this.listeners.set(type, [])
    }
    this.listeners.get(type).push(listener)
  }

  emit(type, event = {}) {
    const list = this.listeners.get(type) ?? []
    for (const listener of list) {
      listener(event)
    }
  }

  open() {
    this.readyState = MockWebSocket.OPEN
    this.emit("open")
  }

  close() {
    this.readyState = MockWebSocket.CLOSED
    this.emit("close", { code: 1000, reason: "Normal" })
  }

  send() {}
}

test("beforeConnect is called and awaited before createSocket", async () => {
  let beforeConnectCalled = false
  let socketCreated = false
  let mockSocket = null

  const manager = createRealtimeConnectionManager({
    beforeConnect: async () => {
      await new Promise((resolve) => setTimeout(resolve, 10))
      beforeConnectCalled = true
    },
    createSocket: () => {
      assert.ok(beforeConnectCalled, "expected beforeConnect to run before createSocket")
      socketCreated = true
      mockSocket = new MockWebSocket()
      return mockSocket
    },
  })

  manager.connect()
  // Wait for beforeConnect promise to resolve
  await new Promise((resolve) => setTimeout(resolve, 30))

  assert.ok(beforeConnectCalled, "beforeConnect should be called")
  assert.ok(socketCreated, "createSocket should be called")
  assert.ok(mockSocket, "socket should be created")
  manager.disconnect()
})

test("beforeConnect failure triggers onConnectError and disconnect status", async () => {
  let connectError = null
  let statusList = []

  const manager = createRealtimeConnectionManager({
    beforeConnect: async () => {
      throw new Error("token expired")
    },
    createSocket: () => {
      assert.fail("createSocket must not be called when beforeConnect fails")
    },
    onConnectError: (err) => {
      connectError = err
    },
    onStatusChange: (status) => {
      statusList.push(status)
    },
  })

  manager.connect()
  await new Promise((resolve) => setTimeout(resolve, 20))

  assert.ok(connectError, "onConnectError should be called")
  assert.equal(connectError.message, "token expired")
  assert.ok(statusList.includes("disconnected"), "status should include disconnected")
  manager.disconnect()
})

test("a stuck CONNECTING socket times out and retries without reloading", async (t) => {
  t.mock.timers.enable({ apis: ["setTimeout", "setInterval"] })
  const sockets = []
  const manager = createRealtimeConnectionManager({
    connectTimeoutMs: 100,
    reconnectBaseDelayMs: 10,
    createSocket: () => { const socket = new MockWebSocket(); sockets.push(socket); return socket },
  })
  t.after(() => manager.disconnect())
  await manager.connect()
  t.mock.timers.tick(100)
  assert.equal(sockets[0].readyState, MockWebSocket.CLOSED)
  t.mock.timers.tick(10)
  assert.equal(sockets.length, 2)
  sockets[1].open()
  t.mock.timers.tick(100)
  assert.equal(manager.getSocket(), sockets[1])
})

test("timed out authentication cannot open a late socket after a newer connection", async (t) => {
  t.mock.timers.enable({ apis: ["setTimeout", "setInterval"] })
  let resolveOld, attempts = 0, created = 0
  const old = new Promise((resolve) => { resolveOld = resolve })
  const manager = createRealtimeConnectionManager({
    connectTimeoutMs: 100,
    reconnectBaseDelayMs: 10,
    beforeConnect: () => ++attempts === 1 ? old : Promise.resolve(),
    createSocket: () => { created++; return new MockWebSocket() },
  })
  t.after(() => manager.disconnect())
  const first = manager.connect()
  t.mock.timers.tick(100)
  t.mock.timers.tick(10)
  await Promise.resolve()
  resolveOld()
  await first
  assert.equal(created, 1)
})

test("late replaced socket close cannot clear the current heartbeat or subscription", async (t) => {
  t.mock.timers.enable({ apis: ["setTimeout", "setInterval"] })
  let closed = 0, pings = 0
  const sockets = []
  const manager = createRealtimeConnectionManager({
    pingIntervalMs: 10,
    createSocket: () => {
      const socket = new MockWebSocket()
      socket.close = () => { socket.readyState = MockWebSocket.CLOSED }
      socket.send = () => { pings++ }
      sockets.push(socket)
      return socket
    },
    onClose: () => { closed++ },
  })
  t.after(() => manager.disconnect())
  await manager.connect()
  sockets[0].open()
  await manager.connect()
  sockets[1].open()
  sockets[0].emit("close", {})
  t.mock.timers.tick(10)
  assert.equal(manager.getSocket(), sockets[1])
  assert.equal(closed, 0)
  assert.equal(pings, 1)
})

test("an error without a close event detaches a connecting socket and retries", async (t) => {
  t.mock.timers.enable({ apis: ["setTimeout", "setInterval"] })
  const sockets = []
  const manager = createRealtimeConnectionManager({
    connectTimeoutMs: 100,
    reconnectBaseDelayMs: 10,
    createSocket: () => {
      const socket = new MockWebSocket()
      socket.close = () => { socket.readyState = MockWebSocket.CLOSING }
      sockets.push(socket)
      return socket
    },
  })
  t.after(() => manager.disconnect())
  await manager.connect()
  sockets[0].emit("error", {})
  assert.equal(manager.getSocket(), null)
  t.mock.timers.tick(10)
  assert.equal(sockets.length, 2)
  sockets[1].open()
  t.mock.timers.tick(100)
  assert.equal(manager.getSocket(), sockets[1])
})

test("becoming ineligible during authentication clears the opening deadline", async (t) => {
  t.mock.timers.enable({ apis: ["setTimeout", "setInterval"] })
  let resolveAuth, allowed = true, errors = 0
  const auth = new Promise((resolve) => { resolveAuth = resolve })
  const manager = createRealtimeConnectionManager({
    connectTimeoutMs: 100,
    canReconnect: () => allowed,
    beforeConnect: () => auth,
    createSocket: () => assert.fail("ineligible connection must not create a socket"),
    onConnectError: () => { errors++ },
  })
  t.after(() => manager.disconnect())
  const connecting = manager.connect()
  allowed = false
  resolveAuth()
  await connecting
  t.mock.timers.tick(200)
  assert.equal(errors, 0)
})

test("disconnect invalidates a pending beforeConnect even across reconnect", async () => {
  let resolveOld
  let attempts = 0, created = 0
  const old = new Promise((resolve) => { resolveOld = resolve })
  const manager = createRealtimeConnectionManager({
    beforeConnect: () => ++attempts === 1 ? old : Promise.resolve(),
    createSocket: () => { created++; return new MockWebSocket() },
  })
  const first = manager.connect()
  manager.disconnect()
  await manager.connect()
  resolveOld()
  await first
  assert.equal(created, 1)
  manager.disconnect()
})

test("late old socket close cannot clear current connection or notify current owner", async () => {
  let closed = 0
  const sockets = []
  const manager = createRealtimeConnectionManager({
    createSocket: () => {
      const socket = new MockWebSocket()
      // Simulate an asynchronously arriving close event after replacement.
      socket.close = () => { socket.readyState = MockWebSocket.CLOSED }
      sockets.push(socket)
      return socket
    },
    onClose: () => { closed++ },
  })
  await manager.connect()
  await manager.connect()
  sockets[0].emit("close", {})
  assert.equal(manager.getSocket(), sockets[1])
  assert.equal(closed, 0)
  manager.disconnect()
})