refactor(utils): extract createReduxConnection

src/utils/redux-extension/createReduxConnection.ts

@@ -0,0 +1,46 @@
+import { Message } from '../types';
+import { ReduxExtension } from './getReduxExtension';
+
+// Original but incomplete type of the redux extension package
+type ConnectResponse = ReturnType<NonNullable<ReduxExtension>['connect']>;
+
+export type Connection = {
+  /** Mark the connection as not initiated, so it can be initiated before using it. */
+  shouldInit?: boolean;
+
+  /** Initiate the connection and add it to the extension connections.
+   *  Should only be executed once in the live time of the connection.
+   */
+  init: ConnectResponse['init'];
+
+  // FIXME https://github.com/reduxjs/redux-devtools/issues/1097
+  /** Add a subscription to the connection.
+   *  The provided listener will be executed when the user interacts with the extension
+   *  with actions like time traveling, importing a state or the likes.
+   *
+   *  @param listener function to be executed when an action is submitted
+   *  @returns function to unsubscribe the applied listener
+   */
+  subscribe: (listener: (message: Message) => void) => (() => void) | undefined;
+
+  /** Send a new action to the connection to display the state change in the extension.
+   *  For example when the value of the store changes.
+   */
+  send: ConnectResponse['send'];
+};
+
+/** Wrapper for creating connections to the redux extension
+ *  Connections are used to display the stores value and value changes within the extension
+ *  as well as reacting to extension actions like time traveling.
+ **/
+export const createReduxConnection = (
+  extension: ReduxExtension | undefined,
+  name: string,
+) => {
+  if (!extension) return undefined;
+  const connection = extension.connect({ name });
+
+  return Object.assign(connection, {
+    shouldInit: true,
+  }) as Connection;
+};

src/utils/useAtomDevtools.ts

@@ -1,8 +1,11 @@
 import { useEffect, useRef } from 'react';
 import { useAtom } from 'jotai/react';
 import type { Atom, WritableAtom } from 'jotai/vanilla';
+import {
+  Connection,
+  createReduxConnection,
+} from './redux-extension/createReduxConnection';
 import { getReduxExtension } from './redux-extension/getReduxExtension';
-import { Message } from './types';
 
 type DevtoolOptions = Parameters<typeof useAtom>[1] & {
   name?: string;
@@ -21,13 +24,7 @@ export function useAtomDevtools<Value, Result>(
 
   const lastValue = useRef(value);
   const isTimeTraveling = useRef(false);
-  const devtools = useRef<
-    ReturnType<
-      NonNullable<typeof window['__REDUX_DEVTOOLS_EXTENSION__']>['connect']
-    > & {
-      shouldInit?: boolean;
-    }
-  >();
+  const devtools = useRef<Connection>();
 
   const atomName = name || anAtom.debugLabel || anAtom.toString();
 
@@ -46,16 +43,9 @@ export function useAtomDevtools<Value, Result>(
       );
     };
 
-    devtools.current = extension.connect({ name: atomName });
+    devtools.current = createReduxConnection(extension, atomName);
 
-    const unsubscribe = (
-      devtools.current as unknown as {
-        // FIXME https://github.com/reduxjs/redux-devtools/issues/1097
-        subscribe: (
-          listener: (message: Message) => void,
-        ) => (() => void) | undefined;
-      }
-    ).subscribe((message) => {
+    const unsubscribe = devtools.current?.subscribe((message) => {
       if (message.type === 'ACTION' && message.payload) {
         try {
           setValueIfWritable(JSON.parse(message.payload));
@@ -95,7 +85,7 @@ export function useAtomDevtools<Value, Result>(
         });
       }
     });
-    devtools.current.shouldInit = true;
+
     return unsubscribe;
   }, [anAtom, extension, atomName, setValue]);
 

src/utils/useAtomsDevtools.ts

@@ -1,7 +1,10 @@
 import { useEffect, useRef } from 'react';
 import { AnyAtom, AnyAtomValue, AtomsSnapshot, Options } from '../types';
+import {
+  Connection,
+  createReduxConnection,
+} from './redux-extension/createReduxConnection';
 import { getReduxExtension } from './redux-extension/getReduxExtension';
-import { Message } from './types';
 import { useAtomsSnapshot } from './useAtomsSnapshot';
 import { useGotoAtomsSnapshot } from './useGotoAtomsSnapshot';
 
@@ -41,13 +44,7 @@ export function useAtomsDevtools(
 
   const isTimeTraveling = useRef(false);
   const isRecording = useRef(true);
-  const devtools = useRef<
-    ReturnType<
-      NonNullable<typeof window['__REDUX_DEVTOOLS_EXTENSION__']>['connect']
-    > & {
-      shouldInit?: boolean;
-    }
-  >();
+  const devtools = useRef<Connection>();
 
   const snapshots = useRef<AtomsSnapshot[]>([]);
 
@@ -63,16 +60,10 @@ export function useAtomsDevtools(
       }
       return snapshot;
     };
-    const connection = extension.connect({ name });
-
-    const devtoolsUnsubscribe = (
-      connection as unknown as {
-        // FIXME https://github.com/reduxjs/redux-devtools/issues/1097
-        subscribe: (
-          listener: (message: Message) => void,
-        ) => (() => void) | undefined;
-      }
-    ).subscribe((message) => {
+
+    devtools.current = createReduxConnection(extension, name);
+
+    const devtoolsUnsubscribe = devtools.current?.subscribe((message) => {
       switch (message.type) {
         case 'DISPATCH':
           switch (message.payload?.type) {
@@ -81,7 +72,7 @@ export function useAtomsDevtools(
               break;
 
             case 'COMMIT':
-              connection.init(getDevtoolsState(getSnapshotAt()));
+              devtools.current?.init(getDevtoolsState(getSnapshotAt()));
               snapshots.current = [];
               break;
 
@@ -98,8 +89,6 @@ export function useAtomsDevtools(
       }
     });
 
-    devtools.current = connection;
-    devtools.current.shouldInit = true;
     return () => {
       extension?.disconnect?.();
       devtoolsUnsubscribe?.();