Merge pull request #22940 from storybookjs/norbert/server-channel-for…

…-addons

Core: Integrate serverChannel into channel, make serverChannel bi-directional and extendable

code/builders/builder-vite/package.json

@@ -43,7 +43,6 @@
   },
   "dependencies": {
     "@storybook/channel-postmessage": "7.1.0-alpha.31",
-    "@storybook/channel-websocket": "7.1.0-alpha.31",
     "@storybook/client-logger": "7.1.0-alpha.31",
     "@storybook/core-common": "7.1.0-alpha.31",
     "@storybook/csf-plugin": "7.1.0-alpha.31",

code/builders/builder-vite/src/codegen-set-addon-channel.ts

@@ -1,17 +1,14 @@
 export async function generateAddonSetupCode() {
   return `
     import { createChannel as createPostMessageChannel } from '@storybook/channel-postmessage';
-    import { createChannel as createWebSocketChannel } from '@storybook/channel-websocket';
     import { addons } from '@storybook/preview-api';
 
     const channel = createPostMessageChannel({ page: 'preview' });
     addons.setChannel(channel);
     window.__STORYBOOK_ADDONS_CHANNEL__ = channel;
     
     if (window.CONFIG_TYPE === 'DEVELOPMENT'){
-      const serverChannel = createWebSocketChannel({});
-      addons.setServerChannel(serverChannel);
-      window.__STORYBOOK_SERVER_CHANNEL__ = serverChannel;
+      window.__STORYBOOK_SERVER_CHANNEL__ = channel;
     }
   `.trim();
 }

code/builders/builder-webpack5/package.json

@@ -59,7 +59,6 @@
     "@storybook/addons": "7.1.0-alpha.31",
     "@storybook/api": "7.1.0-alpha.31",
     "@storybook/channel-postmessage": "7.1.0-alpha.31",
-    "@storybook/channel-websocket": "7.1.0-alpha.31",
     "@storybook/channels": "7.1.0-alpha.31",
     "@storybook/client-api": "7.1.0-alpha.31",
     "@storybook/client-logger": "7.1.0-alpha.31",

code/builders/builder-webpack5/templates/virtualModuleModernEntry.js.handlebars

@@ -2,7 +2,6 @@ import { global } from '@storybook/global';
 
 import { ClientApi, PreviewWeb, addons, composeConfigs } from '@storybook/preview-api';
 import { createChannel as createPostMessageChannel } from '@storybook/channel-postmessage';
-import { createChannel as createWebSocketChannel } from '@storybook/channel-websocket';
 
 import { importFn } from './{{storiesFilename}}';
 
@@ -13,9 +12,7 @@ const channel = createPostMessageChannel({ page: 'preview' });
 addons.setChannel(channel);
 
 if (global.CONFIG_TYPE === 'DEVELOPMENT'){
-  const serverChannel = createWebSocketChannel({});
-  addons.setServerChannel(serverChannel);
-  window.__STORYBOOK_SERVER_CHANNEL__ = serverChannel;
+  window.__STORYBOOK_SERVER_CHANNEL__ = channel;
 }
 
 const preview = new PreviewWeb();

code/frameworks/html-vite/package.json

@@ -50,7 +50,6 @@
     "@storybook/addons": "7.1.0-alpha.31",
     "@storybook/builder-vite": "7.1.0-alpha.31",
     "@storybook/channel-postmessage": "7.1.0-alpha.31",
-    "@storybook/channel-websocket": "7.1.0-alpha.31",
     "@storybook/client-api": "7.1.0-alpha.31",
     "@storybook/core-server": "7.1.0-alpha.31",
     "@storybook/html": "7.1.0-alpha.31",

code/lib/channel-postmessage/package.json

@@ -43,6 +43,7 @@
     "prep": "../../../scripts/prepare/bundle.ts"
   },
   "dependencies": {
+    "@storybook/channel-websocket": "7.1.0-alpha.31",
     "@storybook/channels": "7.1.0-alpha.31",
     "@storybook/client-logger": "7.1.0-alpha.31",
     "@storybook/core-events": "7.1.0-alpha.31",

code/lib/channel-postmessage/src/index.ts

@@ -2,12 +2,13 @@ import { global } from '@storybook/global';
 import * as EVENTS from '@storybook/core-events';
 import { Channel } from '@storybook/channels';
 import type { ChannelHandler, ChannelEvent, ChannelTransport } from '@storybook/channels';
+import { WebsocketTransport } from '@storybook/channel-websocket';
 import { logger, pretty } from '@storybook/client-logger';
 import { isJSON, parse, stringify } from 'telejson';
 import qs from 'qs';
 import invariant from 'tiny-invariant';
 
-const { document, location } = global;
+const { document, location, CONFIG_TYPE } = global;
 
 interface Config {
   page: 'manager' | 'preview';
@@ -289,8 +290,17 @@ const getEventSourceUrl = (event: MessageEvent) => {
  * Creates a channel which communicates with an iframe or child window.
  */
 export function createChannel({ page }: Config): Channel {
-  const transport = new PostmsgTransport({ page });
-  return new Channel({ transport });
+  const transports: ChannelTransport[] = [new PostmsgTransport({ page })];
+
+  if (CONFIG_TYPE === 'DEVELOPMENT') {
+    const protocol = window.location.protocol === 'http:' ? 'ws' : 'wss';
+    const { hostname, port } = window.location;
+    const channelUrl = `${protocol}://${hostname}:${port}/storybook-server-channel`;
+
+    transports.push(new WebsocketTransport({ url: channelUrl, onError: () => {} }));
+  }
+
+  return new Channel({ transports });
 }
 
 // backwards compat with builder-vite

code/lib/channel-postmessage/src/typings.d.ts

@@ -1 +1,2 @@
 declare var CHANNEL_OPTIONS: any;
+declare var CONFIG_TYPE: 'DEVELOPMENT' | 'PRODUCTION';

code/lib/channel-websocket/src/index.ts

@@ -1,10 +1,12 @@
 import { global } from '@storybook/global';
 import { Channel } from '@storybook/channels';
-import type { ChannelHandler } from '@storybook/channels';
+import type { ChannelHandler, ChannelTransport } from '@storybook/channels';
 import { logger } from '@storybook/client-logger';
 import { isJSON, parse, stringify } from 'telejson';
 import invariant from 'tiny-invariant';
 
+const { CONFIG_TYPE } = global;
+
 const { WebSocket } = global;
 
 type OnError = (message: Event) => void;
@@ -80,15 +82,23 @@ export function createChannel({
   async = false,
   onError = (err) => logger.warn(err),
 }: CreateChannelArgs) {
-  let channelUrl = url;
-  if (!channelUrl) {
+  const transports: ChannelTransport[] = [];
+
+  if (url) {
+    transports.push(new WebsocketTransport({ url, onError }));
+  }
+
+  const isUrlServerChannel = !!url?.includes('storybook-server-channel');
+
+  if (CONFIG_TYPE === 'DEVELOPMENT' && isUrlServerChannel === false) {
     const protocol = window.location.protocol === 'http:' ? 'ws' : 'wss';
     const { hostname, port } = window.location;
-    channelUrl = `${protocol}://${hostname}:${port}/storybook-server-channel`;
+    const channelUrl = `${protocol}://${hostname}:${port}/storybook-server-channel`;
+
+    transports.push(new WebsocketTransport({ url: channelUrl, onError: () => {} }));
   }
 
-  const transport = new WebsocketTransport({ url: channelUrl, onError });
-  return new Channel({ transport, async });
+  return new Channel({ transports, async });
 }
 
 // backwards compat with builder-vite

code/lib/channel-websocket/src/typings.d.ts

@@ -1 +1,2 @@
 declare module 'json-fn';
+declare var CONFIG_TYPE: 'DEVELOPMENT' | 'PRODUCTION';

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

@@ -19,7 +19,7 @@ describe('Channel', () => {
     });
 
     it('should not set transport if not passed as an argument', () => {
-      channel = new Channel();
+      channel = new Channel({});
       expect(channel.hasTransport).toBeFalsy();
     });
 
@@ -29,7 +29,7 @@ describe('Channel', () => {
     });
 
     it('should set isAsync to false as default value', () => {
-      channel = new Channel();
+      channel = new Channel({});
       expect(channel.isAsync).toBeFalsy();
     });
 
@@ -104,8 +104,11 @@ describe('Channel', () => {
         listenerOutputData = data;
       });
       const sendSpy = jest.fn();
-      // @ts-expect-error (Converted from ts-ignore)
-      channel.transport.send = sendSpy;
+      // @ts-expect-error (access private property for testing purposes)
+      channel.transports.forEach((t) => {
+        // eslint-disable-next-line no-param-reassign
+        t.send = sendSpy;
+      });
       channel.emit(eventName, ...listenerInputData);
       expect(listenerOutputData).toEqual(listenerInputData);
       expect(sendSpy.mock.calls[0][1]).toEqual({ depth: 1 });

code/lib/channels/src/index.ts

@@ -21,10 +21,20 @@ interface EventsKeyValue {
   [key: string]: Listener[];
 }
 
-interface ChannelArgs {
+type ChannelArgs = ChannelArgsSingle | ChannelArgsMulti;
+interface ChannelArgsSingle {
   transport?: ChannelTransport;
   async?: boolean;
 }
+interface ChannelArgsMulti {
+  transports: ChannelTransport[];
+  async?: boolean;
+}
+
+const isMulti = (args: ChannelArgs): args is ChannelArgsMulti => {
+  // @ts-expect-error (we guard against this right here)
+  return args.transports !== undefined;
+};
 
 const generateRandomId = () => {
   // generates a random 13 character string
@@ -40,18 +50,30 @@ export class Channel {
 
   private data: Record<string, any> = {};
 
-  private readonly transport: ChannelTransport | undefined = undefined;
+  private readonly transports: ChannelTransport[] = [];
+
+  constructor(input: ChannelArgsMulti);
+  constructor(input: ChannelArgsSingle);
+  constructor(input: ChannelArgs = {}) {
+    this.isAsync = input.async || false;
 
-  constructor({ transport, async = false }: ChannelArgs = {}) {
-    this.isAsync = async;
-    if (transport) {
-      this.transport = transport;
-      this.transport.setHandler((event) => this.handleEvent(event));
+    if (isMulti(input)) {
+      this.transports = input.transports || [];
+
+      this.transports.forEach((t) => {
+        t.setHandler((event) => this.handleEvent(event));
+      });
+    } else {
+      this.transports = input.transport ? [input.transport] : [];
     }
+
+    this.transports.forEach((t) => {
+      t.setHandler((event) => this.handleEvent(event));
+    });
   }
 
   get hasTransport() {
-    return !!this.transport;
+    return this.transports.length > 0;
   }
 
   addListener(eventName: string, listener: Listener) {
@@ -67,9 +89,9 @@ export class Channel {
     }
 
     const handler = () => {
-      if (this.transport) {
-        this.transport.send(event, options);
-      }
+      this.transports.forEach((t) => {
+        t.send(event, options);
+      });
       this.handleEvent(event);
     };
 

code/lib/core-server/package.json

@@ -57,6 +57,7 @@
     "@aw-web-design/x-default-browser": "1.4.126",
     "@discoveryjs/json-ext": "^0.5.3",
     "@storybook/builder-manager": "7.1.0-alpha.31",
+    "@storybook/channels": "7.1.0-alpha.31",
     "@storybook/core-common": "7.1.0-alpha.31",
     "@storybook/core-events": "7.1.0-alpha.31",
     "@storybook/csf": "^0.1.0",

code/lib/core-server/src/dev-server.ts

@@ -29,7 +29,10 @@ export async function storybookDevServer(options: Options) {
     options.presets.apply<CoreConfig>('core'),
   ]);
 
-  const serverChannel = getServerChannel(server);
+  const serverChannel = await options.presets.apply(
+    'experimental_serverChannel',
+    getServerChannel(server)
+  );
 
   let indexError: Error;
   // try get index generator, if failed, send telemetry without storyCount, then rethrow the error

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

@@ -1,32 +1,56 @@
 import WebSocket, { WebSocketServer } from 'ws';
-import { stringify } from 'telejson';
+import { isJSON, parse, stringify } from 'telejson';
+import type { ChannelHandler } from '@storybook/channels';
+import { Channel } from '@storybook/channels';
 
 type Server = ConstructorParameters<typeof WebSocketServer>[0]['server'];
 
-export class ServerChannel {
-  webSocketServer: WebSocketServer;
+/**
+ * This class represents a channel transport that allows for a one-to-many relationship between the server and clients.
+ * Unlike other channels such as the postmessage and websocket channel implementations, this channel will receive from many clients and any events emitted will be sent out to all connected clients.
+ */
+export class ServerChannelTransport {
+  private socket: WebSocketServer;
+
+  private handler?: ChannelHandler;
 
   constructor(server: Server) {
-    this.webSocketServer = new WebSocketServer({ noServer: true });
+    this.socket = new WebSocketServer({ noServer: true });
 
     server.on('upgrade', (request, socket, head) => {
       if (request.url === '/storybook-server-channel') {
-        this.webSocketServer.handleUpgrade(request, socket, head, (ws) => {
-          this.webSocketServer.emit('connection', ws, request);
+        this.socket.handleUpgrade(request, socket, head, (ws) => {
+          this.socket.emit('connection', ws, request);
         });
       }
     });
+    this.socket.on('connection', (wss) => {
+      wss.on('message', (raw) => {
+        const data = raw.toString();
+        const event = typeof data === 'string' && isJSON(data) ? parse(data) : data;
+        this.handler(event);
+      });
+    });
+  }
+
+  setHandler(handler: ChannelHandler) {
+    this.handler = handler;
   }
 
-  emit(type: string, args: any = []) {
-    const event = { type, args };
+  send(event: any) {
     const data = stringify(event, { maxDepth: 15, allowFunction: true });
-    Array.from(this.webSocketServer.clients)
+
+    Array.from(this.socket.clients)
       .filter((c) => c.readyState === WebSocket.OPEN)
       .forEach((client) => client.send(data));
   }
 }
 
 export function getServerChannel(server: Server) {
-  return new ServerChannel(server);
+  const transports = [new ServerChannelTransport(server)];
+
+  return new Channel({ transports, async: true });
 }
+
+// for backwards compatibility
+export type ServerChannel = ReturnType<typeof getServerChannel>;

code/lib/manager-api/src/lib/addons.ts

@@ -36,6 +36,9 @@ export class AddonStore {
 
   private channel: Channel | undefined;
 
+  /**
+   * @deprecated will be removed in 8.0
+   */
   private serverChannel: Channel | undefined;
 
   private promise: any;
@@ -51,6 +54,9 @@ export class AddonStore {
     return this.channel;
   };
 
+  /**
+   * @deprecated will be removed in 8.0, use getChannel instead
+   */
   getServerChannel = (): Channel => {
     if (!this.serverChannel) {
       throw new Error('Accessing non-existent serverChannel');
@@ -63,13 +69,19 @@ export class AddonStore {
 
   hasChannel = (): boolean => !!this.channel;
 
+  /**
+   * @deprecated will be removed in 8.0, please use the normal channel instead
+   */
   hasServerChannel = (): boolean => !!this.serverChannel;
 
   setChannel = (channel: Channel): void => {
     this.channel = channel;
     this.resolve();
   };
 
+  /**
+   * @deprecated will be removed in 8.0, please use the normal channel instead
+   */
   setServerChannel = (channel: Channel): void => {
     this.serverChannel = channel;
   };

code/lib/manager-api/src/modules/channel.ts

@@ -74,7 +74,6 @@ export const init: ModuleFn<SubAPI, SubState> = ({ provider }) => {
       }
       provider.channel.emit(type, data, ...args);
     },
-
     collapseAll: () => {
       api.emit(STORIES_COLLAPSE_ALL, {});
     },