Add optional listener options to track message type subscriptions

SSHari · SSHari · commit d81cad892fec · 2023-04-23T00:57:09.000-04:00
diff --git a/sandpack-client/src/clients/base.ts b/sandpack-client/src/clients/base.ts
@@ -6,6 +6,8 @@ import type {
   ListenerFunction,
   SandpackMessage,
   UnsubscribeFunction,
+  SandpackMessageType,
+  ListenerOptions,
 } from "../types";
 
 export class SandpackClient {
@@ -59,7 +61,10 @@ export class SandpackClient {
     throw Error("Method not implemented");
   }
 
-  public listen(_listener: ListenerFunction): UnsubscribeFunction {
+  public listen<Type extends SandpackMessageType = SandpackMessageType>(
+    _listener: ListenerFunction<Type>,
+    _opts?: ListenerOptions<Type>
+  ): UnsubscribeFunction {
     throw Error("Method not implemented");
   }
 }
diff --git a/sandpack-client/src/clients/node/index.ts b/sandpack-client/src/clients/node/index.ts
@@ -13,6 +13,8 @@ import type {
   ListenerFunction,
   SandboxSetup,
   UnsubscribeFunction,
+  SandpackMessageType,
+  ListenerOptions,
 } from "../..";
 import { nullthrows } from "../..";
 import { createError } from "../..";
@@ -434,8 +436,11 @@ export class SandpackNode extends SandpackClient {
     }
   }
 
-  public listen(listener: ListenerFunction): UnsubscribeFunction {
-    return this.emitter.listener(listener);
+  public listen<Type extends SandpackMessageType = SandpackMessageType>(
+    listener: ListenerFunction<Type>,
+    _opts?: ListenerOptions<Type>
+  ): UnsubscribeFunction {
+    return this.emitter.listener(listener as any as ListenerFunction);
   }
 
   public destroy(): void {
diff --git a/sandpack-client/src/clients/runtime/index.ts b/sandpack-client/src/clients/runtime/index.ts
@@ -13,6 +13,8 @@ import type {
   SandpackBundlerFiles,
   SandpackError,
   UnsubscribeFunction,
+  SandpackMessageType,
+  ListenerOptions,
 } from "../../types";
 import { SandpackLogLevel } from "../../types";
 import {
@@ -265,8 +267,13 @@ export class SandpackRuntime extends SandpackClient {
     this.iframeProtocol.dispatch(message);
   }
 
-  public listen(listener: ListenerFunction): UnsubscribeFunction {
-    return this.iframeProtocol.channelListen(listener);
+  public listen<Type extends SandpackMessageType = SandpackMessageType>(
+    listener: ListenerFunction<Type>,
+    _opts?: ListenerOptions<Type>
+  ): UnsubscribeFunction {
+    return this.iframeProtocol.channelListen(
+      listener as any as ListenerFunction
+    );
   }
 
   /**
diff --git a/sandpack-client/src/clients/static/index.ts b/sandpack-client/src/clients/static/index.ts
@@ -4,8 +4,11 @@ import { PreviewController } from "static-browser-server";
 
 import type {
   ClientOptions,
+  ListenerOptions,
   ListenerFunction,
   SandboxSetup,
+  SandpackMessageType,
+  SandpackMessageByType,
   UnsubscribeFunction,
 } from "../..";
 import { SandpackClient } from "../base";
@@ -23,6 +26,7 @@ export class SandpackStatic extends SandpackClient {
   private emitter: EventEmitter;
   private previewController: PreviewController;
   private files: Map<string, string | Uint8Array> = new Map();
+  private registeredMessageTypes: Map<SandpackMessageType, number> = new Map();
 
   public iframe!: HTMLIFrameElement;
   public selector!: string;
@@ -54,9 +58,12 @@ export class SandpackStatic extends SandpackClient {
               content,
               options.externalResources
             );
-            content = this.injectScriptIntoHead(content, {
-              script: consoleHook,
-            });
+
+            if (this.registeredMessageTypes.has("console")) {
+              content = this.injectScriptIntoHead(content, {
+                script: consoleHook,
+              });
+            }
           } catch (err) {
             console.error("Runtime injection failed", err);
           }
@@ -223,8 +230,37 @@ export class SandpackStatic extends SandpackClient {
     }
   }
 
-  public listen(listener: ListenerFunction): UnsubscribeFunction {
-    return this.emitter.listener(listener);
+  public listen<Type extends SandpackMessageType = SandpackMessageType>(
+    listener: ListenerFunction<Type>,
+    opts?: ListenerOptions<Type>
+  ): UnsubscribeFunction {
+    opts?.messageTypes?.forEach((type) => {
+      const count = this.registeredMessageTypes.get(type) ?? 0;
+      this.registeredMessageTypes.set(type, count + 1);
+    });
+
+    const unsubscribe = this.emitter.listener((message) => {
+      if (
+        !opts?.messageTypes ||
+        opts.messageTypes.includes(message.type as Type)
+      ) {
+        listener(message as SandpackMessageByType<Type>);
+      }
+    });
+
+    return () => {
+      unsubscribe();
+
+      opts?.messageTypes?.forEach((type) => {
+        const count = this.registeredMessageTypes.get(type)!;
+
+        if (count === 1) {
+          this.registeredMessageTypes.delete(type);
+        } else {
+          this.registeredMessageTypes.set(type, count - 1);
+        }
+      });
+    };
   }
 
   public destroy(): void {
diff --git a/sandpack-client/src/types.ts b/sandpack-client/src/types.ts
@@ -174,13 +174,27 @@ export type ClientStatus =
   | "idle";
 
 export type SandpackMessage = SandpackRuntimeMessage | SandpackNodeMessage;
+export type SandpackMessageType = SandpackMessage["type"];
+export type SandpackMessageByType<Type extends SandpackMessageType> = Extract<
+  SandpackMessage,
+  { type: Type }
+>;
 
-export type ListenerFunction = (msg: SandpackMessage) => void;
+export type ListenerFunction<
+  Type extends SandpackMessageType = SandpackMessageType
+> = (msg: SandpackMessageByType<Type>) => void;
 export type UnsubscribeFunction = () => void;
 
-export type Listen = (
-  listener: ListenerFunction,
-  clientId?: string
+export type ListenerOptions<
+  Type extends SandpackMessageType = SandpackMessageType
+> = {
+  messageTypes?: Array<Type>;
+};
+
+export type Listen = <Type extends SandpackMessageType = SandpackMessageType>(
+  listener: ListenerFunction<Type>,
+  clientId?: string,
+  opts?: ListenerOptions<Type>
 ) => UnsubscribeFunction;
 export type Dispatch = (msg: SandpackMessage, clientId?: string) => void;
 
diff --git a/sandpack-react/src/components/Console/useSandpackConsole.ts b/sandpack-react/src/components/Console/useSandpackConsole.ts
@@ -29,52 +29,56 @@ export const useSandpackConsole = ({
   const { listen } = useSandpack();
 
   React.useEffect(() => {
-    const unsubscribe = listen((message) => {
-      if (resetOnPreviewRestart && message.type === "start") {
-        setLogs([]);
-      } else if (message.type === "console" && message.codesandbox) {
-        const payloadLog = Array.isArray(message.log)
-          ? message.log
-          : [message.log];
+    const unsubscribe = listen(
+      (message) => {
+        if (resetOnPreviewRestart && message.type === "start") {
+          setLogs([]);
+        } else if (message.type === "console" && message.codesandbox) {
+          const payloadLog = Array.isArray(message.log)
+            ? message.log
+            : [message.log];
 
-        if (payloadLog.find(({ method }) => method === "clear")) {
-          return setLogs([CLEAR_LOG]);
-        }
+          if (payloadLog.find(({ method }) => method === "clear")) {
+            return setLogs([CLEAR_LOG]);
+          }
 
-        const logsMessages = showSyntaxError
-          ? payloadLog
-          : payloadLog.filter((messageItem) => {
-              const messagesWithoutSyntaxErrors =
-                messageItem?.data?.filter?.((dataItem) => {
-                  if (typeof dataItem !== "string") return true;
+          const logsMessages = showSyntaxError
+            ? payloadLog
+            : payloadLog.filter((messageItem) => {
+                const messagesWithoutSyntaxErrors =
+                  messageItem?.data?.filter?.((dataItem) => {
+                    if (typeof dataItem !== "string") return true;
 
-                  const matches = SYNTAX_ERROR_PATTERN.filter((lookFor) =>
-                    dataItem.startsWith(lookFor)
-                  );
+                    const matches = SYNTAX_ERROR_PATTERN.filter((lookFor) =>
+                      dataItem.startsWith(lookFor)
+                    );
 
-                  return matches.length === 0;
-                }) ?? [];
+                    return matches.length === 0;
+                  }) ?? [];
 
-              return messagesWithoutSyntaxErrors.length > 0;
-            });
+                return messagesWithoutSyntaxErrors.length > 0;
+              });
 
-        if (!logsMessages) return;
+          if (!logsMessages) return;
 
-        setLogs((prev) => {
-          const messages = [...prev, ...logsMessages].filter(
-            (value, index, self) => {
-              return index === self.findIndex((s) => s.id === value.id);
-            }
-          );
+          setLogs((prev) => {
+            const messages = [...prev, ...logsMessages].filter(
+              (value, index, self) => {
+                return index === self.findIndex((s) => s.id === value.id);
+              }
+            );
 
-          while (messages.length > maxMessageCount) {
-            messages.shift();
-          }
+            while (messages.length > maxMessageCount) {
+              messages.shift();
+            }
 
-          return messages;
-        });
-      }
-    }, clientId);
+            return messages;
+          });
+        }
+      },
+      clientId,
+      { messageTypes: ["console", "start"] }
+    );
 
     return unsubscribe;
   }, [showSyntaxError, maxMessageCount, clientId, resetOnPreviewRestart]);
diff --git a/sandpack-react/src/contexts/utils/useClient.ts b/sandpack-react/src/contexts/utils/useClient.ts
@@ -6,6 +6,8 @@ import type {
   SandpackMessage,
   UnsubscribeFunction,
   SandpackClient,
+  SandpackMessageType,
+  ListenerOptions,
 } from "@codesandbox/sandpack-client";
 import {
   loadSandpackClient,
@@ -49,9 +51,10 @@ export interface UseClientOperations {
     clientPropsOverride?: ClientPropsOverride
   ) => Promise<void>;
   registerReactDevTools: (value: ReactDevToolsMode) => void;
-  addListener: (
+  addListener: <Type extends SandpackMessageType = SandpackMessageType>(
     listener: ListenerFunction,
-    clientId?: string
+    clientId?: string,
+    opts?: ListenerOptions<Type>
   ) => UnsubscribeFunction;
   dispatchMessage: (message: SandpackMessage, clientId?: string) => void;
   lazyAnchorRef: React.RefObject<HTMLDivElement>;
@@ -61,7 +64,10 @@ export interface UseClientOperations {
     Record<string, Record<string, UnsubscribeFunction>>
   >;
   queuedListenersRef: React.MutableRefObject<
-    Record<string, Record<string, ListenerFunction>>
+    Record<
+      string,
+      Record<string, { listener: ListenerFunction; opts?: ListenerOptions }>
+    >
   >;
 }
 
@@ -106,7 +112,10 @@ export const useClient: UseClient = (
   >({});
   const unsubscribe = useRef<() => void | undefined>();
   const queuedListeners = useRef<
-    Record<string, Record<string, ListenerFunction>>
+    Record<
+      string,
+      Record<string, { listener: ListenerFunction; opts?: ListenerOptions }>
+    >
   >({ global: {} });
   const debounceHook = useRef<number | undefined>();
   const loadingScreenRegisteredRef = useRef<boolean>(true);
@@ -175,8 +184,9 @@ export const useClient: UseClient = (
        */
       if (queuedListeners.current[clientId]) {
         Object.keys(queuedListeners.current[clientId]).forEach((listenerId) => {
-          const listener = queuedListeners.current[clientId][listenerId];
-          const unsubscribe = client.listen(listener) as () => void;
+          const { listener, opts } =
+            queuedListeners.current[clientId][listenerId];
+          const unsubscribe = client.listen(listener, opts) as () => void;
           unsubscribeClientListeners.current[clientId][listenerId] =
             unsubscribe;
         });
@@ -189,8 +199,8 @@ export const useClient: UseClient = (
        * Register global listeners
        */
       const globalListeners = Object.entries(queuedListeners.current.global);
-      globalListeners.forEach(([listenerId, listener]) => {
-        const unsubscribe = client.listen(listener) as () => void;
+      globalListeners.forEach(([listenerId, { listener, opts }]) => {
+        const unsubscribe = client.listen(listener, opts) as () => void;
         unsubscribeClientListeners.current[clientId][listenerId] = unsubscribe;
 
         /**
@@ -394,13 +404,17 @@ export const useClient: UseClient = (
     }
   };
 
-  const addListener = (
-    listener: ListenerFunction,
-    clientId?: string
+  const addListener = <Type extends SandpackMessageType = SandpackMessageType>(
+    listener: ListenerFunction<Type>,
+    clientId?: string,
+    opts?: ListenerOptions<Type>
   ): UnsubscribeFunction => {
     if (clientId) {
       if (clients.current[clientId]) {
-        const unsubscribeListener = clients.current[clientId].listen(listener);
+        const unsubscribeListener = clients.current[clientId].listen(
+          listener,
+          opts
+        );
 
         return unsubscribeListener;
       } else {
@@ -415,7 +429,10 @@ export const useClient: UseClient = (
         unsubscribeClientListeners.current[clientId] =
           unsubscribeClientListeners.current[clientId] || {};
 
-        queuedListeners.current[clientId][listenerId] = listener;
+        queuedListeners.current[clientId][listenerId] = {
+          listener: listener as any as ListenerFunction,
+          opts: opts as any as ListenerOptions,
+        };
 
         const unsubscribeListener = (): void => {
           if (queuedListeners.current[clientId][listenerId]) {
@@ -439,7 +456,10 @@ export const useClient: UseClient = (
     } else {
       // Push to the **global** queue
       const listenerId = generateRandomId();
-      queuedListeners.current.global[listenerId] = listener;
+      queuedListeners.current.global[listenerId] = {
+        listener: listener as any as ListenerFunction,
+        opts: opts as any as ListenerOptions,
+      };
 
       // Add to the current clients
       const clientsList = Object.values(clients.current);
diff --git a/sandpack-react/src/types.ts b/sandpack-react/src/types.ts

Original file line number	Diff line number	Diff line change
`@@ -6,6 +6,8 @@ import type {`
`6`	`6`	`ListenerFunction,`
`7`	`7`	`SandpackMessage,`
`8`	`8`	`UnsubscribeFunction,`
	`9`	`+ SandpackMessageType,`
	`10`	`+ ListenerOptions,`
`9`	`11`	`} from "../types";`
`10`	`12`
`11`	`13`	`export class SandpackClient {`
`@@ -59,7 +61,10 @@ export class SandpackClient {`
`59`	`61`	`throw Error("Method not implemented");`
`60`	`62`	`}`
`61`	`63`
`62`		`- public listen(_listener: ListenerFunction): UnsubscribeFunction {`
	`64`	`+ public listen<Type extends SandpackMessageType = SandpackMessageType>(`
	`65`	`+ _listener: ListenerFunction<Type>,`
	`66`	`+ _opts?: ListenerOptions<Type>`
	`67`	`+ ): UnsubscribeFunction {`
`63`	`68`	`throw Error("Method not implemented");`
`64`	`69`	`}`
`65`	`70`	`}`