Merge pull request #30288 from storybookjs/valentin/implement-connect…

…ion-closed-notification

Core: Add connection timeout notification

code/core/package.json

@@ -275,6 +275,7 @@
   },
   "dependencies": {
     "@storybook/csf": "0.1.12",
+    "@storybook/theming": "workspace:*",
     "better-opn": "^3.0.2",
     "browser-assert": "^1.2.1",
     "esbuild": "^0.18.0 || ^0.19.0 || ^0.20.0 || ^0.21.0 || ^0.22.0 || ^0.23.0 || ^0.24.0",

code/core/scripts/entries.ts

@@ -66,6 +66,6 @@ export const getFinals = (cwd: string) => {
 
   return [
     //
-    define('src/manager/runtime.ts', ['browser'], false),
+    define('src/manager/runtime.tsx', ['browser'], false),
   ];
 };

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

@@ -12,7 +12,7 @@ const MockedWebsocket = vi.hoisted(() => {
 
     onerror: (e: any) => void;
 
-    onclose: () => void;
+    onclose: (event: any) => void;
 
     constructor(url: string) {
       this.onopen = vi.fn();
@@ -293,11 +293,16 @@ describe('WebsocketTransport', () => {
     });
 
     transport.setHandler(handler);
-    MockedWebsocket.ref.current.onclose();
+    MockedWebsocket.ref.current.onclose({ code: 1000, reason: 'test' });
 
     expect(handler.mock.calls[0][0]).toMatchInlineSnapshot(`
       {
-        "args": [],
+        "args": [
+          {
+            "code": 1000,
+            "reason": "test",
+          },
+        ],
         "from": "preview",
         "type": "channelWSDisconnect",
       }

code/core/src/channels/index.ts

@@ -13,7 +13,7 @@ export * from './main';
 export default Channel;
 
 export { PostMessageTransport } from './postmessage';
-export { WebsocketTransport } from './websocket';
+export { WebsocketTransport, HEARTBEAT_INTERVAL, HEARTBEAT_MAX_LATENCY } from './websocket';
 
 type Options = Config & {
   extraTransports?: ChannelTransport[];

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

@@ -17,6 +17,9 @@ interface WebsocketTransportArgs extends Partial<Config> {
   onError: OnError;
 }
 
+export const HEARTBEAT_INTERVAL = 15000;
+export const HEARTBEAT_MAX_LATENCY = 5000;
+
 export class WebsocketTransport implements ChannelTransport {
   private buffer: string[] = [];
 
@@ -26,25 +29,48 @@ export class WebsocketTransport implements ChannelTransport {
 
   private isReady = false;
 
+  private isClosed = false;
+
+  private pingTimeout: number | NodeJS.Timeout = 0;
+
+  private heartbeat() {
+    clearTimeout(this.pingTimeout);
+
+    this.pingTimeout = setTimeout(() => {
+      this.socket.close(3008, 'timeout');
+    }, HEARTBEAT_INTERVAL + HEARTBEAT_MAX_LATENCY);
+  }
+
   constructor({ url, onError, page }: WebsocketTransportArgs) {
     this.socket = new WebSocket(url);
     this.socket.onopen = () => {
       this.isReady = true;
+      this.heartbeat();
       this.flush();
     };
     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);
+      if (event.type === 'ping') {
+        this.heartbeat();
+        this.send({ type: 'pong' });
+      }
     };
     this.socket.onerror = (e) => {
       if (onError) {
         onError(e);
       }
     };
-    this.socket.onclose = () => {
+    this.socket.onclose = (ev) => {
       invariant(this.handler, 'WebsocketTransport handler should be set');
-      this.handler({ type: EVENTS.CHANNEL_WS_DISCONNECT, args: [], from: page || 'preview' });
+      this.handler({
+        type: EVENTS.CHANNEL_WS_DISCONNECT,
+        args: [{ reason: ev.reason, code: ev.code }],
+        from: page || 'preview',
+      });
+      this.isClosed = true;
+      clearTimeout(this.pingTimeout);
     };
   }
 
@@ -53,10 +79,12 @@ export class WebsocketTransport implements ChannelTransport {
   }
 
   send(event: any) {
-    if (!this.isReady) {
-      this.sendLater(event);
-    } else {
-      this.sendNow(event);
+    if (!this.isClosed) {
+      if (!this.isReady) {
+        this.sendLater(event);
+      } else {
+        this.sendNow(event);
+      }
     }
   }
 

code/core/src/core-server/utils/get-server-channel.ts

@@ -1,5 +1,5 @@
 import type { ChannelHandler } from '@storybook/core/channels';
-import { Channel } from '@storybook/core/channels';
+import { Channel, HEARTBEAT_INTERVAL } from '@storybook/core/channels';
 
 import { isJSON, parse, stringify } from 'telejson';
 import WebSocket, { WebSocketServer } from 'ws';
@@ -17,7 +17,15 @@ export class ServerChannelTransport {
 
   private handler?: ChannelHandler;
 
+  isAlive = false;
+
+  private heartbeat() {
+    this.isAlive = true;
+  }
+
   constructor(server: Server) {
+    this.heartbeat = this.heartbeat.bind(this);
+
     this.socket = new WebSocketServer({ noServer: true });
 
     server.on('upgrade', (request, socket, head) => {
@@ -28,14 +36,42 @@ export class ServerChannelTransport {
       }
     });
     this.socket.on('connection', (wss) => {
+      this.isAlive = true;
       wss.on('message', (raw) => {
         const data = raw.toString();
         const event =
           typeof data === 'string' && isJSON(data)
             ? parse(data, { allowFunction: false, allowClass: false })
             : data;
         this.handler?.(event);
+        if (event.type === 'pong') {
+          this.heartbeat();
+        }
+      });
+    });
+
+    const interval = setInterval(() => {
+      this.socket.clients.forEach((ws) => {
+        if (this.isAlive === false) {
+          return ws.terminate();
+        }
+
+        this.isAlive = false;
+        this.send({ type: 'ping' });
+      });
+    }, HEARTBEAT_INTERVAL);
+
+    this.socket.on('close', function close() {
+      clearInterval(interval);
+    });
+
+    process.on('SIGTERM', () => {
+      this.socket.clients.forEach((client) => {
+        if (client.readyState === WebSocket.OPEN) {
+          client.close(1001, 'Server is shutting down');
+        }
       });
+      this.socket.close(() => process.exit(0));
     });
   }
 

code/core/src/manager/globals/exports.ts

@@ -743,18 +743,24 @@ export default {
   '@storybook/core/theming/create': ['create', 'themes'],
   'storybook/internal/channels': [
     'Channel',
+    'HEARTBEAT_INTERVAL',
+    'HEARTBEAT_MAX_LATENCY',
     'PostMessageTransport',
     'WebsocketTransport',
     'createBrowserChannel',
   ],
   '@storybook/channels': [
     'Channel',
+    'HEARTBEAT_INTERVAL',
+    'HEARTBEAT_MAX_LATENCY',
     'PostMessageTransport',
     'WebsocketTransport',
     'createBrowserChannel',
   ],
   '@storybook/core/channels': [
     'Channel',
+    'HEARTBEAT_INTERVAL',
+    'HEARTBEAT_MAX_LATENCY',
     'PostMessageTransport',
     'WebsocketTransport',
     'createBrowserChannel',

code/core/src/manager/runtime.ts → code/core/src/manager/runtime.tsx

@@ -1,20 +1,28 @@
+import React from 'react';
+
 import type { Channel } from '@storybook/core/channels';
 import { createBrowserChannel } from '@storybook/core/channels';
 import type { Addon_Config, Addon_Types } from '@storybook/core/types';
 import { global } from '@storybook/global';
+import { FailedIcon } from '@storybook/icons';
+import { color } from '@storybook/theming';
 
-import { CHANNEL_CREATED } from '@storybook/core/core-events';
-import type { AddonStore } from '@storybook/core/manager-api';
+import { CHANNEL_CREATED, CHANNEL_WS_DISCONNECT } from '@storybook/core/core-events';
+import type { API, AddonStore } from '@storybook/core/manager-api';
 import { addons } from '@storybook/core/manager-api';
 
 import { renderStorybookUI } from './index';
 import Provider from './provider';
 
+const WS_DISCONNECTED_NOTIFICATION_ID = 'CORE/WS_DISCONNECTED';
+
 class ReactProvider extends Provider {
   addons: AddonStore;
 
   channel: Channel;
 
+  wsDisconnected = false;
+
   constructor() {
     super();
 
@@ -37,8 +45,23 @@ class ReactProvider extends Provider {
     return this.addons.getConfig();
   }
 
-  handleAPI(api: unknown) {
+  handleAPI(api: API) {
     this.addons.loadAddons(api);
+
+    this.channel.on(CHANNEL_WS_DISCONNECT, (ev) => {
+      const TIMEOUT_CODE = 3008;
+      this.wsDisconnected = true;
+
+      api.addNotification({
+        id: WS_DISCONNECTED_NOTIFICATION_ID,
+        content: {
+          headline: ev.code === TIMEOUT_CODE ? 'Server timed out' : 'Connection lost',
+          subHeadline: 'Please restart your Storybook server and reload the page',
+        },
+        icon: <FailedIcon color={color.negative} />,
+        link: undefined,
+      });
+    });
   }
 }
 

code/renderers/web-components/src/docs/web-components-properties.test.ts

@@ -33,36 +33,35 @@ const runWebComponentsAnalyzer = (inputPath: string) => {
   return output;
 };
 
-describe('web-components component properties', () => {
+vi.mock('lit', () => ({ default: {} }));
+vi.mock('lit/directive-helpers.js', () => ({ default: {} }));
+
+describe('web-components component properties', { timeout: 15000 }, () => {
   // we need to mock lit and dynamically require custom-elements
   // because lit is distributed as ESM not CJS
   // https://github.com/Polymer/lit-html/issues/516
-  vi.mock('lit', () => ({ default: {} }));
-  vi.mock('lit/directive-helpers.js', () => ({ default: {} }));
 
   const fixturesDir = join(__dirname, '__testfixtures__');
-  readdirSync(fixturesDir, { withFileTypes: true }).forEach((testEntry) => {
+  const testEntries = readdirSync(fixturesDir, { withFileTypes: true });
+
+  it.each(testEntries)('$name', async (testEntry) => {
     if (testEntry.isDirectory()) {
       const testDir = join(fixturesDir, testEntry.name);
       const testFile = readdirSync(testDir).find((fileName) => inputRegExp.test(fileName));
       if (testFile) {
-        it(`${testEntry.name}`, async () => {
-          const inputPath = join(testDir, testFile);
+        const inputPath = join(testDir, testFile);
 
-          // snapshot the output of wca
-          const customElementsJson = runWebComponentsAnalyzer(inputPath);
-          const customElements = JSON.parse(customElementsJson);
-          customElements.tags.forEach((tag: any) => {
-            tag.path = 'dummy-path-to-component';
-          });
-          await expect(customElements).toMatchFileSnapshot(
-            join(testDir, 'custom-elements.snapshot')
-          );
-
-          // snapshot the properties
-          const properties = extractArgTypesFromElements('input', customElements);
-          await expect(properties).toMatchFileSnapshot(join(testDir, 'properties.snapshot'));
+        // snapshot the output of wca
+        const customElementsJson = runWebComponentsAnalyzer(inputPath);
+        const customElements = JSON.parse(customElementsJson);
+        customElements.tags.forEach((tag: any) => {
+          tag.path = 'dummy-path-to-component';
         });
+        await expect(customElements).toMatchFileSnapshot(join(testDir, 'custom-elements.snapshot'));
+
+        // snapshot the properties
+        const properties = extractArgTypesFromElements('input', customElements);
+        await expect(properties).toMatchFileSnapshot(join(testDir, 'properties.snapshot'));
       }
     }
   });

code/yarn.lock

@@ -6258,6 +6258,7 @@ __metadata:
     "@storybook/docs-mdx": "npm:4.0.0-next.1"
     "@storybook/global": "npm:^5.0.0"
     "@storybook/icons": "npm:^1.2.12"
+    "@storybook/theming": "workspace:*"
     "@tanstack/react-virtual": "npm:^3.3.0"
     "@testing-library/react": "npm:^14.0.0"
     "@types/cross-spawn": "npm:^6.0.2"