fix(relay-run): buffer early ws frames in wsToWebSocketLike

The message listener was attached only when the mux consumer called onMessage — which
happens inside attach(), AFTER the async mTLS-registry lookup. On loopback the agent's FIRST
heartbeat ping arrived in that gap and ws dropped it (no listener), so the agent's heartbeat
died on the first miss at 15s and the tunnel flapped forever. Attach the listener at
construction and buffer early frames (and an early close) until the consumer wires up.
This commit is contained in:
Yaojia Wang
2026-07-07 04:59:53 +02:00
parent 89678c7949
commit c1c837c54f

View File

@@ -82,8 +82,33 @@ function toUint8(data: unknown): Uint8Array | null {
return null
}
/** Adapt a live `ws` WebSocket to `WebSocketLike` (binary frames only). */
/**
* Adapt a live `ws` WebSocket to `WebSocketLike` (binary frames only).
*
* The `ws.on('message'|'close')` listeners are attached NOW (at construction), not when the consumer
* later calls `onMessage`/`onClose`. The composition root builds this synchronously on the socket's
* 'connection' event, but the mux consumer only wires its handler AFTER an async mTLS-registry lookup
* (the bridge). Any frame that arrives in that gap — notably the agent's FIRST heartbeat ping, sent
* immediately on open — would otherwise be dropped by `ws` (no listener yet), starving the heartbeat
* and flapping the tunnel every 15 s. We buffer early frames (and an early close) and flush on wire-up.
*/
export function wsToWebSocketLike(ws: WsWebSocket): WebSocketLike {
let onMessage: ((data: Uint8Array) => void) | null = null
let onClose: (() => void) | null = null
const backlog: Uint8Array[] = []
let closedEarly = false
ws.on('message', (data: unknown) => {
const bytes = toUint8(data)
if (bytes === null) return
if (onMessage === null) backlog.push(bytes)
else onMessage(bytes)
})
ws.on('close', () => {
if (onClose === null) closedEarly = true
else onClose()
})
return {
send(data: Uint8Array) {
ws.send(data, { binary: true })
@@ -92,13 +117,12 @@ export function wsToWebSocketLike(ws: WsWebSocket): WebSocketLike {
ws.close(code)
},
onMessage(handler) {
ws.on('message', (data: unknown) => {
const bytes = toUint8(data)
if (bytes !== null) handler(bytes)
})
onMessage = handler
while (backlog.length > 0) handler(backlog.shift()!)
},
onClose(handler) {
ws.on('close', () => handler())
onClose = handler
if (closedEarly) handler()
},
}
}