Newer
Older
navi-1 / webclient / tests / unit / composables / useWebSocket.test.js
import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest'
import { createPinia, setActivePinia } from 'pinia'
import { useWebSocket } from '@/composables/useWebSocket.js'
import { wsSend } from '@/composables/wsBridge.js'
import { useChatStore } from '@/stores/chat.js'

const toastMock = vi.hoisted(() => ({
  show: vi.fn(), info: vi.fn(), success: vi.fn(), warning: vi.fn(), error: vi.fn(),
}))
vi.mock('gnexus-ui-kit/vue', () => ({ useToast: () => toastMock }))

class MockWebSocket extends EventTarget {
  // The composable compares readyState against WebSocket.OPEN/CLOSED and the
  // watchdog against WebSocket.CONNECTING. Stubbing the global without these
  // makes every constant undefined, so any readyState matches any constant.
  static CONNECTING = 0
  static OPEN = 1
  static CLOSING = 2
  static CLOSED = 3

  constructor(url) {
    super()
    this.url = url
    this.readyState = WebSocket.CONNECTING
    this.sent = []
    MockWebSocket.last = this
  }

  send(data) { this.sent.push(data) }
  close() {
    this.readyState = WebSocket.CLOSED
    this.dispatchEvent(new Event('close'))
  }

  _open() {
    this.readyState = WebSocket.OPEN
    this.dispatchEvent(new Event('open'))
  }

  _msg(data) {
    this.dispatchEvent(new MessageEvent('message', { data: JSON.stringify(data) }))
  }

  // Close with an application code — what the server does for 4003/4004.
  _closeWith(code, reason = '') {
    this.readyState = WebSocket.CLOSED
    const event = new Event('close')
    event.code = code
    event.reason = reason
    this.dispatchEvent(event)
  }
}

// Stub global WebSocket before each test
beforeEach(() => {
  vi.stubGlobal('WebSocket', MockWebSocket)
  // happy-dom does not provide localStorage — getWsUrl reads it for the API token
  vi.stubGlobal('localStorage', {
    getItem: vi.fn(() => null),
    setItem: vi.fn(),
    removeItem: vi.fn(),
  })
  setActivePinia(createPinia())
})

afterEach(() => {
  vi.unstubAllGlobals()
})

describe('useWebSocket', () => {
  it('connect opens socket', () => {
    const ws = useWebSocket()
    ws.connect('s1')
    expect(MockWebSocket.last).toBeTruthy()
    expect(MockWebSocket.last.url).toContain('/ws/sessions/s1')
  })

  it('connected flag after open', () => {
    const ws = useWebSocket()
    ws.connect('s1')
    MockWebSocket.last._open()
    expect(ws.connected.value).toBe(true)
    expect(ws.reconnecting.value).toBe(false)
  })

  it('send delivers JSON when open', () => {
    const ws = useWebSocket()
    ws.connect('s1')
    MockWebSocket.last._open()
    ws.send({ type: 'message', content: 'hi' })
    expect(MockWebSocket.last.sent).toHaveLength(1)
    expect(JSON.parse(MockWebSocket.last.sent[0])).toEqual({ type: 'message', content: 'hi' })
  })

  it('registers its send on the bridge for cards outside the provide chain', () => {
    const ws = useWebSocket()
    ws.connect('s1')
    MockWebSocket.last._open()

    const payload = { type: 'skill_proposal', proposal_id: 'p1', action: 'accept' }
    expect(wsSend(payload)).toBe(true)
    expect(JSON.parse(MockWebSocket.last.sent[0])).toEqual(payload)
  })

  it('disconnect sets destroyed and clears state', () => {
    const ws = useWebSocket()
    ws.connect('s1')
    MockWebSocket.last._open()
    ws.disconnect()
    expect(ws.connected.value).toBe(false)
  })

  it('dispatches stream_start to chat store', () => {
    const ws = useWebSocket()
    ws.connect('s1')
    MockWebSocket.last._open()
    MockWebSocket.last._msg({ type: 'stream_start' })
    // No crash = dispatched; deeper assertions need chat store spy
    expect(ws.connected.value).toBe(true)
  })

  it('dispatches session_sync', () => {
    const ws = useWebSocket()
    ws.connect('s1')
    MockWebSocket.last._open()
    MockWebSocket.last._msg({ type: 'session_sync' })
    expect(ws.connected.value).toBe(true)
  })

  it('dispatches task_update and message_queued to the chat store', () => {
    const ws = useWebSocket()
    const chat = useChatStore()
    ws.connect('s1')
    MockWebSocket.last._open()
    MockWebSocket.last._msg({ type: 'task_update', task_id: 'bt-7', tool: 'peer', status: 'running' })
    expect(chat.backgroundTasks['bt-7'].tool).toBe('peer')

    chat.awaitingStream = true
    MockWebSocket.last._msg({ type: 'message_queued', position: 1, queue_len: 1, max: 5 })
    expect(chat.awaitingStream).toBe(false)
  })

  it('passes the display rank through on stream_stopped', () => {
    const ws = useWebSocket()
    const chat = useChatStore()
    const stopped = vi.spyOn(chat, 'onStreamStopped')
    ws.connect('s1')
    MockWebSocket.last._open()

    MockWebSocket.last._msg({ type: 'stream_stopped', message_index: 9 })

    expect(stopped).toHaveBeenCalledWith({ type: 'stream_stopped', message_index: 9 })
    ws.disconnect()
  })

  it('dispatches terminal_opened to the chat store', () => {
    const ws = useWebSocket()
    const chat = useChatStore()
    ws.connect('s1')
    MockWebSocket.last._open()
    MockWebSocket.last._msg({ type: 'terminal_opened', terminal_name: 'ops', description: '', pid: 3, background: true })
    expect(chat.terminals['ops']).toMatchObject({ name: 'ops', closed: false })
  })

  it('closes a pending handshake when the session changes', () => {
    const ws = useWebSocket()
    ws.connect('s1')
    const first = MockWebSocket.last
    expect(first.url).toContain('/ws/sessions/s1')

    // The handshake to s1 has not completed — the client has already moved on.
    ws.connect('s2')

    expect(first.readyState).toBe(WebSocket.CLOSED)
    expect(first.onmessage).toBe(null)
    expect(MockWebSocket.last).not.toBe(first)
    expect(MockWebSocket.last.url).toContain('/ws/sessions/s2')
  })

  it('a late event from a replaced socket cannot reach the transcript', () => {
    const ws = useWebSocket()
    const chat = useChatStore()
    ws.connect('s1')
    const first = MockWebSocket.last
    first._open()

    ws.connect('s2')
    // Driving the old socket after the switch must be a no-op.
    first._open()
    first._msg({ type: 'stream_delta', delta: 'from the previous session' })
    first._msg({ type: 'stream_start' })

    expect(chat.messages).toHaveLength(0)
    expect(chat.streaming).toBe(false)
  })

  it('does not open a second socket for a handshake already in flight', () => {
    const ws = useWebSocket()
    ws.connect('s1')
    const first = MockWebSocket.last
    ws.connect('s1')
    expect(MockWebSocket.last).toBe(first)
  })

  it('reconnect schedules on close', () => {
    vi.useFakeTimers()
    const ws = useWebSocket()
    ws.connect('s1')
    MockWebSocket.last._open()
    MockWebSocket.last.close()
    expect(ws.reconnecting.value).toBe(true)
    vi.advanceTimersByTime(3000)
    // Should have created a new MockWebSocket instance
    expect(MockWebSocket.last.readyState).toBe(WebSocket.CONNECTING)
    vi.useRealTimers()
  })
})

describe('useWebSocket — task_update toasts', () => {
  beforeEach(() => {
    vi.clearAllMocks()
  })

  it('toasts start for a new task and finish on terminal status', () => {
    const ws = useWebSocket()
    const chat = useChatStore()
    chat.currentId = 's1'
    ws.connect('s1')
    MockWebSocket.last._open()
    MockWebSocket.last._msg({ type: 'task_update', task_id: 'bt-9', tool: 'terminal', status: 'running', session_id: 's1' })
    expect(toastMock.info).toHaveBeenCalledTimes(1)
    expect(toastMock.info.mock.calls[0][0]).toMatchObject({ title: 'Background task started' })
    // repeated running publishes must not re-toast
    MockWebSocket.last._msg({ type: 'task_update', task_id: 'bt-9', tool: 'terminal', status: 'running', session_id: 's1' })
    expect(toastMock.info).toHaveBeenCalledTimes(1)
    MockWebSocket.last._msg({ type: 'task_update', task_id: 'bt-9', tool: 'terminal', status: 'completed', result_preview: '', session_id: 's1' })
    expect(toastMock.success).toHaveBeenCalledTimes(1)
  })

  it('does not toast a late running update for a finished task', () => {
    const ws = useWebSocket()
    const chat = useChatStore()
    chat.currentId = 's1'
    ws.connect('s1')
    MockWebSocket.last._open()
    MockWebSocket.last._msg({ type: 'task_update', task_id: 'bt-8', tool: 'peer', status: 'running', session_id: 's1' })
    toastMock.info.mockClear()
    MockWebSocket.last._msg({ type: 'task_update', task_id: 'bt-8', tool: 'peer', status: 'completed', session_id: 's1' })
    expect(toastMock.success).toHaveBeenCalledTimes(1)
    toastMock.success.mockClear()
    MockWebSocket.last._msg({ type: 'task_update', task_id: 'bt-8', tool: 'peer', status: 'running', session_id: 's1' })
    expect(toastMock.info).not.toHaveBeenCalled()
  })

  it('skips toasts for other sessions', () => {
    const ws = useWebSocket()
    const chat = useChatStore()
    chat.currentId = 's1'
    ws.connect('s1')
    MockWebSocket.last._open()
    MockWebSocket.last._msg({ type: 'task_update', task_id: 'bt-o', tool: 'terminal', status: 'running', session_id: 's-other' })
    expect(toastMock.info).not.toHaveBeenCalled()
  })
})

describe('useWebSocket — liveness and recovery', () => {
  it('pings the server while the connection sits idle', () => {
    vi.useFakeTimers()
    const ws = useWebSocket()
    ws.connect('s1')
    const sock = MockWebSocket.last
    sock._open()
    vi.advanceTimersByTime(20000)
    expect(sock.sent.map(d => JSON.parse(d).type)).toContain('ping')
    ws.disconnect()
    vi.useRealTimers()
  })

  it('answers to a server heartbeat prove the socket is alive', () => {
    vi.useFakeTimers()
    const ws = useWebSocket()
    ws.connect('s1')
    const sock = MockWebSocket.last
    sock._open()
    // Heartbeats arrive every 20 s — nothing should be considered stale.
    for (let i = 0; i < 4; i++) {
      vi.advanceTimersByTime(20000)
      sock._msg({ type: 'heartbeat' })
    }
    expect(ws.connected.value).toBe(true)
    expect(MockWebSocket.last).toBe(sock)
    ws.disconnect()
    vi.useRealTimers()
  })

  it('queues a message instead of writing to a silent socket', () => {
    vi.useFakeTimers()
    const ws = useWebSocket()
    ws.connect('s1')
    const first = MockWebSocket.last
    first._open()
    // No frames for two heartbeat windows: half-open, writes would vanish.
    vi.advanceTimersByTime(46000)
    expect(ws.send({ type: 'message', content: 'hi' })).toBe('queued')
    expect(first.sent.some(d => d.includes('"message"'))).toBe(false)
    expect(MockWebSocket.last).not.toBe(first)
    ws.disconnect()
    vi.useRealTimers()
  })

  it('rejects a message instead of dropping the oldest queued one', () => {
    vi.useFakeTimers()
    const ws = useWebSocket()
    ws.connect('s1')
    const first = MockWebSocket.last
    first._open()
    // Park the socket in the stale state so every send goes to the queue.
    vi.advanceTimersByTime(46000)
    for (let i = 0; i < 20; i++) {
      expect(ws.send({ type: 'message', content: `m${i}` })).toBe('queued')
    }
    expect(ws.send({ type: 'message', content: 'overflow' })).toBe('rejected')

    // Reconnect: the twenty kept payloads flush, in order, and the rejected
    // one is the only one missing — the oldest were not silently evicted.
    MockWebSocket.last._open()
    const sent = MockWebSocket.last.sent.map(d => JSON.parse(d))
      .filter(p => p.type === 'message')
      .map(p => p.content)
    expect(sent).toEqual(Array.from({ length: 20 }, (_, i) => `m${i}`))
    ws.disconnect()
    vi.useRealTimers()
  })

  it('drops the stream that died with the socket and syncs history', () => {
    vi.useFakeTimers()
    const chat = useChatStore()
    const ws = useWebSocket()
    ws.connect('s1')
    const first = MockWebSocket.last
    first._open()

    chat.onStreamStart()
    chat.onStreamDelta('half an answer')
    expect(chat.streaming).toBe(true)

    first.close()
    vi.advanceTimersByTime(1000)
    MockWebSocket.last._open()
    expect(chat.streaming).toBe(false)
    expect(chat.messages).toHaveLength(0)

    // The run finished while we were away — session_sync must get through now.
    // It uses the cheap meta check, not a full history reload.
    const check = vi.spyOn(chat, 'checkForUpdates').mockResolvedValue(false)
    MockWebSocket.last._msg({ type: 'session_sync' })
    expect(check).toHaveBeenCalledWith('s1')
    ws.disconnect()
    vi.useRealTimers()
  })

  it('ignores session_sync while the session is still loading', () => {
    vi.useFakeTimers()
    const chat = useChatStore()
    const ws = useWebSocket()
    ws.connect('s1')
    MockWebSocket.last._open()
    chat.loading = true

    const check = vi.spyOn(chat, 'checkForUpdates').mockResolvedValue(false)
    MockWebSocket.last._msg({ type: 'session_sync' })
    expect(check).not.toHaveBeenCalled()
    chat.loading = false
    ws.disconnect()
    vi.useRealTimers()
  })

  it('reconnects immediately when the network comes back', () => {
    const ws = useWebSocket()
    ws.connect('s1')
    const first = MockWebSocket.last
    first._open()

    window.dispatchEvent(new Event('online'))

    expect(MockWebSocket.last).not.toBe(first)
    ws.disconnect()
  })
})

describe('useWebSocket — delivery acknowledgement', () => {
  it('clears the deadline when the server accepts the message', () => {
    vi.useFakeTimers()
    const chat = useChatStore()
    const ws = useWebSocket()
    chat.appendUserMessage('hi', [], [], { clientKey: 'c1' })
    ws.connect('s1')
    MockWebSocket.last._open()

    expect(ws.send({ type: 'message', content: 'hi', client_id: 'c1' })).toBe('sent')
    MockWebSocket.last._msg({ type: 'message_accepted', client_id: 'c1' })
    vi.advanceTimersByTime(20000)

    expect(chat.messages[0].failed).toBeUndefined()
    ws.disconnect()
    vi.useRealTimers()
  })

  it('a queued message counts as accepted too', () => {
    vi.useFakeTimers()
    const chat = useChatStore()
    const ws = useWebSocket()
    chat.appendUserMessage('hi', [], [], { clientKey: 'c2' })
    ws.connect('s1')
    MockWebSocket.last._open()

    ws.send({ type: 'message', content: 'hi', client_id: 'c2' })
    MockWebSocket.last._msg({ type: 'message_queued', position: 1, queue_len: 1, max: 5, client_id: 'c2' })
    vi.advanceTimersByTime(20000)

    expect(chat.messages[0].failed).toBeUndefined()
    ws.disconnect()
    vi.useRealTimers()
  })

  it('marks the bubble when no ack arrives within the deadline', () => {
    vi.useFakeTimers()
    const chat = useChatStore()
    const ws = useWebSocket()
    chat.appendUserMessage('hi', [], [], { clientKey: 'c3' })
    ws.connect('s1')
    MockWebSocket.last._open()

    ws.send({ type: 'message', content: 'hi', client_id: 'c3' })
    vi.advanceTimersByTime(15000)

    expect(chat.messages[0].failed).toBe(true)
    expect(chat.messages[0].failReason).toBe('unconfirmed')
    ws.disconnect()
    vi.useRealTimers()
  })

  it('arming starts at the write, not at the queueing', () => {
    vi.useFakeTimers()
    const chat = useChatStore()
    const ws = useWebSocket()
    chat.appendUserMessage('hi', [], [], { clientKey: 'c4' })
    ws.connect('s1')
    const first = MockWebSocket.last
    first._open()
    // Park the socket in the stale state: the payload waits in the queue.
    vi.advanceTimersByTime(46000)
    expect(ws.send({ type: 'message', content: 'hi', client_id: 'c4' })).toBe('queued')

    // No write happened yet — there is nothing for the server to acknowledge.
    vi.advanceTimersByTime(20000)
    expect(chat.messages[0].failed).toBeUndefined()

    MockWebSocket.last._open()
    vi.advanceTimersByTime(15000)
    expect(chat.messages[0].failed).toBe(true)
    ws.disconnect()
    vi.useRealTimers()
  })

  it('waits out an active run instead of calling the message lost', () => {
    vi.useFakeTimers()
    const chat = useChatStore()
    const ws = useWebSocket()
    chat.appendUserMessage('hi', [], [], { clientKey: 'c6' })
    ws.connect('s1')
    MockWebSocket.last._open()

    // A run is streaming: the server will not read this frame until it ends,
    // so the first deadline expiring proves nothing.
    chat.onStreamStart()
    ws.send({ type: 'message', content: 'hi', client_id: 'c6' })
    vi.advanceTimersByTime(15000)
    expect(chat.messages[0].failed).toBeUndefined()

    // The run ends without ever reading the frame — now silence is a loss.
    chat.onStreamEnd({ message_index: 1 })
    vi.advanceTimersByTime(15000)
    expect(chat.messages[0].failed).toBe(true)
    expect(chat.messages[0].failReason).toBe('unconfirmed')
    ws.disconnect()
    vi.useRealTimers()
  })

  it('a message the server dropped stops looking queued', () => {
    const chat = useChatStore()
    const ws = useWebSocket()
    chat.appendUserMessage('hi', [], [], { queued: true, clientKey: 'c5' })
    ws.connect('s1')
    MockWebSocket.last._open()

    MockWebSocket.last._msg({ type: 'message_dropped', client_ids: ['c5', 'someone-else'] })

    expect(chat.messages[0].queued).toBe(false)
    expect(chat.messages[0].failed).toBe(true)
    expect(chat.messages[0].failReason).toBe('dropped')
    ws.disconnect()
  })
})

describe('useWebSocket — fatal close codes', () => {
  it('4004 stops reconnecting and says the session is gone', () => {
    vi.useFakeTimers()
    const ws = useWebSocket()
    ws.connect('s1')
    const first = MockWebSocket.last
    first._open()

    first._closeWith(4004, 'Session not found')
    vi.advanceTimersByTime(60000)

    expect(ws.fatalClose.value).toEqual({ code: 4004, reason: 'Session not found' })
    expect(ws.reconnecting.value).toBe(false)
    expect(MockWebSocket.last).toBe(first)  // no new socket was opened
    expect(ws.send({ type: 'message', content: 'hi' })).toBe('rejected')
    ws.disconnect()
    vi.useRealTimers()
  })

  it('4003 stops reconnecting and offers the login', () => {
    vi.useFakeTimers()
    const ws = useWebSocket()
    ws.connect('s1')
    const first = MockWebSocket.last
    first._open()

    first._closeWith(4003, 'Access denied')
    vi.advanceTimersByTime(60000)

    expect(ws.fatalClose.value).toEqual({ code: 4003, reason: 'Access denied' })
    expect(ws.reconnecting.value).toBe(false)
    expect(MockWebSocket.last).toBe(first)
    ws.disconnect()
    vi.useRealTimers()
  })

  it('a new connect clears the fatal state', () => {
    const ws = useWebSocket()
    ws.connect('s1')
    MockWebSocket.last._open()
    MockWebSocket.last._closeWith(4004, 'Session not found')
    expect(ws.fatalClose.value).not.toBe(null)

    ws.connect('s2')
    expect(ws.fatalClose.value).toBe(null)
    expect(MockWebSocket.last.url).toContain('/ws/sessions/s2')
    ws.disconnect()
  })

  it('a plain close still reconnects', () => {
    vi.useFakeTimers()
    const ws = useWebSocket()
    ws.connect('s1')
    const first = MockWebSocket.last
    first._open()

    first._closeWith(1006, '')
    vi.advanceTimersByTime(3000)

    expect(ws.fatalClose.value).toBe(null)
    expect(MockWebSocket.last).not.toBe(first)
    ws.disconnect()
    vi.useRealTimers()
  })
})