File size: 2,471 Bytes
1f21206 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 | import { afterEach, describe, expect, it, mock, spyOn } from 'bun:test'
import type { ServerWebSocket } from 'bun'
import {
__resetWebSocketHandlerStateForTests,
closeSessionConnection,
getActiveSessionIds,
handleWebSocket,
type WebSocketData,
} from '../ws/handler.js'
import { conversationService } from '../services/conversationService.js'
import { computerUseApprovalService } from '../services/computerUseApprovalService.js'
function makeClientSocket(sessionId: string) {
const sent: string[] = []
return {
data: {
sessionId,
connectedAt: Date.now(),
channel: 'client',
sdkToken: null,
serverPort: 0,
serverHost: '127.0.0.1',
},
send: mock((payload: string) => {
sent.push(payload)
}),
close: mock(() => {}),
sent,
} as unknown as ServerWebSocket<WebSocketData> & { sent: string[] }
}
describe('WebSocket handler session isolation', () => {
afterEach(() => {
__resetWebSocketHandlerStateForTests()
mock.restore()
})
it('ignores stale disconnects from an older socket for the same session', () => {
const sessionId = `duplicate-${crypto.randomUUID()}`
const first = makeClientSocket(sessionId)
const second = makeClientSocket(sessionId)
const clearCallbacks = spyOn(conversationService, 'clearOutputCallbacks')
const cancelComputerUse = spyOn(computerUseApprovalService, 'cancelSession')
handleWebSocket.open(first)
handleWebSocket.open(second)
clearCallbacks.mockClear()
cancelComputerUse.mockClear()
handleWebSocket.close(first, 1000, 'stale tab closed')
expect(getActiveSessionIds()).toContain(sessionId)
expect(clearCallbacks).not.toHaveBeenCalled()
expect(cancelComputerUse).not.toHaveBeenCalled()
})
it('closes and removes an active client socket when a session is deleted', () => {
const sessionId = `delete-${crypto.randomUUID()}`
const ws = makeClientSocket(sessionId)
const clearCallbacks = spyOn(conversationService, 'clearOutputCallbacks')
const cancelComputerUse = spyOn(computerUseApprovalService, 'cancelSession')
handleWebSocket.open(ws)
expect(closeSessionConnection(sessionId, 'session deleted')).toBe(true)
expect(getActiveSessionIds()).not.toContain(sessionId)
expect(ws.close).toHaveBeenCalledWith(1000, 'session deleted')
expect(clearCallbacks).toHaveBeenCalledWith(sessionId)
expect(cancelComputerUse).toHaveBeenCalledWith(sessionId)
})
})
|