Core: Reset WebsocketTransport heartbeat on every message, not just ping

[valentinpalkovic]

Manual testing

N?A

What I did

Related to #35400, which fixes a "Server timed out" manager disconnect during large @storybook/addon-vitest runs. That PR's core-channel fix reorders WebsocketTransport.onmessage to check for ping and reset the heartbeat before calling the (potentially slow, when the channel is busy) channel handler, instead of after.

This PR implements that same underlying fix directly against next (code/core/src/channels/websocket/index.ts, which doesn't yet have #35400's change), but generalizes it one step further: the heartbeat now resets on every incoming message, not only on ping.

Why: any message arriving over the socket is itself proof the connection is alive - the server is actively sending data and it's reaching this tab. Gating the reset on the literal ping type leaves a gap: if the message backlog in front of the next ping is deep enough (many status-update events, each requiring real synchronous work to handle), the ping can still be delayed past the timeout window even while the connection is constantly, verifiably receiving traffic. Resetting on every message closes that gap directly, since forward progress through the backlog - of any kind - is what actually matters, not whether a ping happens to be interspersed in it.

this.socket.onmessage = ({ data }) => {
  const event = ...;
  invariant(this.handler, 'WebsocketTransport handler should be set');

  this.heartbeat(); // reset on any message, before the handler runs

  if (event.type === 'ping') {
    this.send({ type: 'pong' });
    return; // pings have no channel listeners
  }

  this.handler(event);
};

Notes for reviewers

• Ports and extends the test coverage from Addon-Vitest: Prevent manager "Server timed out" during large test runs #35400's websocket/index.test.ts (which doesn't exist on next yet): timeout-close, ping reset + pong reply, ping not forwarded to the handler, non-ping forwarded, heartbeat reset survives a throwing handler - plus two new tests specifically for the generalized behavior: a non-ping message also resets the heartbeat, and does so even if the handler throws while processing it.
• This only touches the shared client-side transport, used by both the manager and preview. Small diff, one file plus its test.
• Doesn't include addon-vitest's STORYBOOK_TEST_STATUS_FLUSH_INTERVAL env var or anything from Addon-Vitest: Prevent manager "Server timed out" during large test runs #35400/Generic heartbeat hardening on top of the Vitest addon timeout fix #35419 - this is meant to be a minimal, independently mergeable improvement to the shared transport.

Testing

• unit tests (7 tests in websocket/index.test.ts, all passing)
• manual verification against a large storybook (same repro as Addon-Vitest: Prevent manager "Server timed out" during large test runs #35400)

yarn nx run core:check passes with no type errors. Formatted with oxfmt; lint has no new warnings beyond what already existed on this file.

Summary by CodeRabbit

Summary by CodeRabbit

• Bug Fixes
  • Improved WebSocket heartbeat monitoring so every incoming message reliably resets the timeout.
  • Internal ping messages are now auto-answered with pong and are no longer delivered to the message handler.
  • ping no longer triggers handler errors; non-ping messages still forward normally (including error propagation).
  • When heartbeat times out, the connection closes with a timeout error.
• Tests
  • Added coverage for heartbeat timeouts, ping/pong behavior, forwarding rules, and handler error handling.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: ec303555-1818-4496-90e7-8c2004178292

📥 Commits

Reviewing files that changed from the base of the PR and between d0969af and 80483a7.

📒 Files selected for processing (1)

• code/core/src/channels/websocket/index.test.ts

🚧 Files skipped from review as they are similar to previous changes (1)

• code/core/src/channels/websocket/index.test.ts

---

📝 Walkthrough

Walkthrough

WebsocketTransport now resets its heartbeat on every incoming message and sends pong for ping messages without calling the channel handler. A new test suite covers timeout closing, ping/pong behavior, message forwarding, and handler error cases.

Changes

WebSocket heartbeat and ping handling

Layer / File(s)	Summary
Onmessage heartbeat reset and ping short-circuit with tests code/core/src/channels/websocket/index.ts, code/core/src/channels/websocket/index.test.ts	Heartbeat reset now runs for every message before event-type checks; ping messages send a pong and return early without invoking the channel handler; non-ping messages invoke the handler as before. New tests cover timeout closing, ping/pong exchange, handler forwarding rules, and heartbeat reset persisting even when the handler throws.

Estimated code review effort: 3 (Moderate) | ~20 minutes

✨ Finishing Touches

📝 Generate docstrings

• Create stacked PR
• Commit on current branch

---

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@code/core/src/channels/websocket/index.test.ts`:
- Line 50: The WebSocket test is stubbing a global with vi.stubGlobal but never
restoring it, which can leak MockWebSocket into other tests. Update the
websocket test setup in index.test.ts so the stubbed WebSocket global is cleaned
up by adding vi.unstubAllGlobals() in the existing afterEach alongside
vi.useRealTimers() and vi.clearAllMocks(), and keep the stub usage centralized
around MockWebSocket so the suite leaves no ambient globals behind.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: 75bc0610-feeb-49a0-8bb1-523bffa36a63

📥 Commits

Reviewing files that changed from the base of the PR and between cbf8e4f and 5775319.

📒 Files selected for processing (2)

• code/core/src/channels/websocket/index.test.ts
• code/core/src/channels/websocket/index.ts

[storybook-app-bot]

Package Benchmarks

^{Commit: 80483a7, ran on 9 July 2026 at 10:56:27 UTC}

No significant changes detected, all good. 👏

[ndelangen]

@valentinpalkovic, this replaces my PR, right?
#35400

Shall I close it?

[valentinpalkovic]

Yes, though it needs some testing. Since you have replicated the error, can you try out the canary release (which I am triggering now)?

[ndelangen]

I was still testing this @Sidnioulz, I'll need to do a follow-up, as this PR alone did not resolve the problem fully