From 57753194229158bb768e5f722cc33a819b39b985 Mon Sep 17 00:00:00 2001 From: Valentin Palkovic Date: Wed, 8 Jul 2026 21:01:12 +0200 Subject: [PATCH 1/3] Core: Reset WebsocketTransport heartbeat on every message, not just ping Fixes the "Server timed out" manager disconnect on large addon-vitest runs (see #35400) at its root, in the shared channel transport rather than only reordering ping handling: any incoming message proves the connection is alive, so the heartbeat now resets before the (possibly slow) channel handler runs for every message, not only for a literal ping. This also closes a gap the narrower ping-only fix leaves open - if a backlog is deep enough that the next ping itself is stuck behind many other queued messages, resetting only on ping doesn't help, since nothing about processing those other messages pushes the deadline out. --- .../core/src/channels/websocket/index.test.ts | 169 ++++++++++++++++++ code/core/src/channels/websocket/index.ts | 15 +- 2 files changed, 182 insertions(+), 2 deletions(-) create mode 100644 code/core/src/channels/websocket/index.test.ts diff --git a/code/core/src/channels/websocket/index.test.ts b/code/core/src/channels/websocket/index.test.ts new file mode 100644 index 000000000000..bf909c7039db --- /dev/null +++ b/code/core/src/channels/websocket/index.test.ts @@ -0,0 +1,169 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; + +import { parse, stringify } from 'telejson'; + +import { HEARTBEAT_INTERVAL, HEARTBEAT_MAX_LATENCY, WebsocketTransport } from './index.ts'; + +const TIMEOUT = HEARTBEAT_INTERVAL + HEARTBEAT_MAX_LATENCY; + +const socketRef = { current: undefined as unknown as MockWebSocket }; + +/** + * Minimal WebSocket stub that keeps outgoing `send` calls separate from incoming messages, so tests + * can assert heartbeat behavior without a pong echoing straight back into the message handler. + */ +class MockWebSocket { + static OPEN = 1; + + onopen?: () => void; + + onmessage?: (event: { data: string }) => void; + + onerror?: (event: unknown) => void; + + onclose?: (event: { code: number; reason: string }) => void; + + readyState = MockWebSocket.OPEN; + + sent: string[] = []; + + closed?: { code: number; reason: string }; + + constructor(public url: string) { + socketRef.current = this; + } + + send(data: string) { + this.sent.push(data); + } + + close(code: number, reason: string) { + this.closed = { code, reason }; + this.onclose?.({ code, reason }); + } + + receive(event: unknown) { + this.onmessage?.({ data: stringify(event) }); + } +} + +vi.stubGlobal('WebSocket', MockWebSocket); + +const createConnectedTransport = () => { + const handler = vi.fn(); + const onError = vi.fn(); + const transport = new WebsocketTransport({ + url: 'ws://localhost:6006', + page: 'manager', + onError, + }); + transport.setHandler(handler); + // onopen marks the transport ready and starts the heartbeat watchdog + socketRef.current.onopen?.(); + return { transport, handler, socket: socketRef.current }; +}; + +const sentEvents = (socket: MockWebSocket) => socket.sent.map((data) => parse(data)); + +describe('WebsocketTransport heartbeat', () => { + beforeEach(() => { + vi.useFakeTimers(); + }); + + afterEach(() => { + vi.useRealTimers(); + vi.clearAllMocks(); + }); + + it('closes with code 3008 when no message is received within the timeout window', () => { + const { socket } = createConnectedTransport(); + + vi.advanceTimersByTime(TIMEOUT - 1); + expect(socket.closed).toBeUndefined(); + + vi.advanceTimersByTime(1); + expect(socket.closed).toEqual({ code: 3008, reason: 'timeout' }); + }); + + it('resets the heartbeat and replies with pong when a ping is received', () => { + const { socket } = createConnectedTransport(); + + vi.advanceTimersByTime(TIMEOUT - 1); + socket.receive({ type: 'ping' }); + + expect(sentEvents(socket)).toContainEqual({ type: 'pong' }); + + // The heartbeat was reset, so another near-full window can elapse without closing + vi.advanceTimersByTime(TIMEOUT - 1); + expect(socket.closed).toBeUndefined(); + + vi.advanceTimersByTime(1); + expect(socket.closed).toEqual({ code: 3008, reason: 'timeout' }); + }); + + it('resets the heartbeat when a non-ping event is received', () => { + const { socket } = createConnectedTransport(); + + vi.advanceTimersByTime(TIMEOUT - 1); + socket.receive({ type: 'test', args: [], from: 'preview' }); + + // Any message is evidence the connection is alive, not just a literal ping - so this arrival + // must reset the watchdog too, even though it's not the `ping` type. + vi.advanceTimersByTime(TIMEOUT - 1); + expect(socket.closed).toBeUndefined(); + + vi.advanceTimersByTime(1); + expect(socket.closed).toEqual({ code: 3008, reason: 'timeout' }); + }); + + it('does not forward ping events to the channel handler', () => { + const { handler, socket } = createConnectedTransport(); + + socket.receive({ type: 'ping' }); + + expect(handler).not.toHaveBeenCalled(); + }); + + it('forwards non-ping events to the channel handler', () => { + const { handler, socket } = createConnectedTransport(); + + socket.receive({ type: 'test', args: [], from: 'preview' }); + + expect(handler).toHaveBeenCalledWith(expect.objectContaining({ type: 'test' })); + }); + + it('resets the heartbeat for pings even when the channel handler throws', () => { + const { transport, socket } = createConnectedTransport(); + + // Simulate an overloaded/failing channel handler. Heartbeat processing must not depend on it, + // otherwise a busy channel could starve the reset and close a healthy connection with code 3008. + transport.setHandler(() => { + throw new Error('handler boom'); + }); + + vi.advanceTimersByTime(TIMEOUT - 1); + expect(() => socket.receive({ type: 'ping' })).not.toThrow(); + expect(sentEvents(socket)).toContainEqual({ type: 'pong' }); + + vi.advanceTimersByTime(TIMEOUT - 1); + expect(socket.closed).toBeUndefined(); + }); + + it('resets the heartbeat for non-ping events even when the channel handler throws', () => { + const { transport, socket } = createConnectedTransport(); + + transport.setHandler(() => { + throw new Error('handler boom'); + }); + + vi.advanceTimersByTime(TIMEOUT - 1); + // The handler throwing is expected to propagate to the caller (the socket's own onmessage + // dispatch), but the heartbeat reset must already have happened before that point. + expect(() => socket.receive({ type: 'test', args: [], from: 'preview' })).toThrow( + 'handler boom' + ); + + vi.advanceTimersByTime(TIMEOUT - 1); + expect(socket.closed).toBeUndefined(); + }); +}); diff --git a/code/core/src/channels/websocket/index.ts b/code/core/src/channels/websocket/index.ts index 8ed1f0110dc5..a821f9127fc1 100644 --- a/code/core/src/channels/websocket/index.ts +++ b/code/core/src/channels/websocket/index.ts @@ -50,11 +50,22 @@ export class WebsocketTransport implements ChannelTransport { this.socket.onmessage = ({ data }) => { const event = typeof data === 'string' && isJSON(data) ? parse(data) : data; invariant(this.handler, 'WebsocketTransport handler should be set'); - this.handler(event); + + // Any message arriving proves the connection is alive, so reset the heartbeat before running + // the handler rather than gating the reset behind a specific event type. The channel handler + // can be slow when the channel is busy (e.g. addon-vitest flushing status updates during a + // large test run), and a backlog of queued messages could otherwise delay resetting the + // heartbeat past the timeout and close a healthy connection - relying on `ping` alone doesn't + // fully fix this if the backlog is deep enough that the ping itself is stuck behind it. + this.heartbeat(); + if (event.type === 'ping') { - this.heartbeat(); + // Pings are internal to the transport and have no channel listeners. this.send({ type: 'pong' }); + return; } + + this.handler(event); }; this.socket.onerror = (e) => { if (onError) { From d0969af315830cb2d82b86c32a6d232ed5394e1a Mon Sep 17 00:00:00 2001 From: Valentin Palkovic Date: Wed, 8 Jul 2026 21:09:01 +0200 Subject: [PATCH 2/3] Apply suggestion from @valentinpalkovic --- code/core/src/channels/websocket/index.ts | 6 ------ 1 file changed, 6 deletions(-) diff --git a/code/core/src/channels/websocket/index.ts b/code/core/src/channels/websocket/index.ts index a821f9127fc1..7226d5e8b260 100644 --- a/code/core/src/channels/websocket/index.ts +++ b/code/core/src/channels/websocket/index.ts @@ -51,12 +51,6 @@ export class WebsocketTransport implements ChannelTransport { const event = typeof data === 'string' && isJSON(data) ? parse(data) : data; invariant(this.handler, 'WebsocketTransport handler should be set'); - // Any message arriving proves the connection is alive, so reset the heartbeat before running - // the handler rather than gating the reset behind a specific event type. The channel handler - // can be slow when the channel is busy (e.g. addon-vitest flushing status updates during a - // large test run), and a backlog of queued messages could otherwise delay resetting the - // heartbeat past the timeout and close a healthy connection - relying on `ping` alone doesn't - // fully fix this if the backlog is deep enough that the ping itself is stuck behind it. this.heartbeat(); if (event.type === 'ping') { From b85ad5fe62bdad4af5e1f69b07b0ec6c6cafc38d Mon Sep 17 00:00:00 2001 From: Norbert de Langen Date: Thu, 9 Jul 2026 12:33:54 +0200 Subject: [PATCH 3/3] Address review: stub WebSocket in beforeEach and unstub in afterEach Move the `WebSocket` global stub out of module scope into `beforeEach` and add `vi.unstubAllGlobals()` to `afterEach`, so the suite no longer leaves an ambient global behind for other test files in the same worker. CodeRabbit's proposed fix (only adding `vi.unstubAllGlobals()` while keeping the stub at module scope) would break the suite: after the first test's cleanup the stub is gone and subsequent tests construct a real WebSocket, so `socketRef` is never populated. Re-stubbing per test keeps all 7 tests passing while complying with the repo's global-stubbing guideline. Co-authored-by: Cursor --- code/core/src/channels/websocket/index.test.ts | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/code/core/src/channels/websocket/index.test.ts b/code/core/src/channels/websocket/index.test.ts index bf909c7039db..85caba875992 100644 --- a/code/core/src/channels/websocket/index.test.ts +++ b/code/core/src/channels/websocket/index.test.ts @@ -47,8 +47,6 @@ class MockWebSocket { } } -vi.stubGlobal('WebSocket', MockWebSocket); - const createConnectedTransport = () => { const handler = vi.fn(); const onError = vi.fn(); @@ -68,11 +66,13 @@ const sentEvents = (socket: MockWebSocket) => socket.sent.map((data) => parse(da describe('WebsocketTransport heartbeat', () => { beforeEach(() => { vi.useFakeTimers(); + vi.stubGlobal('WebSocket', MockWebSocket); }); afterEach(() => { vi.useRealTimers(); vi.clearAllMocks(); + vi.unstubAllGlobals(); }); it('closes with code 3008 when no message is received within the timeout window', () => {