feat: DH-21093: Make usePersistentState available to all panels (#2595)

Removed `extendState` from GL because we did not use it (and don't in
DHE). It also tripped me up because it mutates the existing object, so
calling it does not trigger a state change for Dashboard which means it
doesn't actually save any extended state on its own.

Author: mattrunyon

packages/dashboard/src/Dashboard.test.tsx

@@ -1,7 +1,22 @@
 import React from 'react';
-import { render, type RenderResult } from '@testing-library/react';
+import {
+  act,
+  render,
+  type RenderResult,
+  screen,
+  waitFor,
+} from '@testing-library/react';
+import userEvent from '@testing-library/user-event';
 import { ApiContext } from '@deephaven/jsapi-bootstrap';
+import { type LayoutManager } from '@deephaven/golden-layout';
 import Dashboard, { type DashboardProps } from './Dashboard';
+import { usePersistentState } from './usePersistentState';
+import { useDashboardPanel } from './layout';
+import {
+  assertIsDashboardPluginProps,
+  type DashboardPluginComponentProps,
+} from './DashboardPlugin';
+import PanelEvent from './PanelEvent';
 
 const mockDispatch = jest.fn();
 jest.mock('react-redux', () => ({
@@ -18,7 +33,7 @@ function makeDashboard({
   layoutConfig,
   layoutSettings,
   onGoldenLayoutChange,
-  onLayoutConfigChange,
+  onLayoutConfigChange = jest.fn(),
 }: DashboardProps = {}): RenderResult {
   return render(
     <ApiContext.Provider value={dh}>
@@ -39,3 +54,92 @@ function makeDashboard({
 it('mounts and unmounts properly', () => {
   makeDashboard();
 });
+
+function TestComponent() {
+  const [state, setState] = usePersistentState('initial', {
+    type: 'test',
+    version: 1,
+  });
+
+  return (
+    <button
+      type="button"
+      onClick={() => {
+        setState('updated');
+      }}
+    >
+      {state}
+    </button>
+  );
+}
+
+TestComponent.displayName = 'TestComponent';
+
+function TestPlugin(props: Partial<DashboardPluginComponentProps>) {
+  assertIsDashboardPluginProps(props);
+
+  useDashboardPanel({
+    componentName: TestComponent.displayName,
+    component: TestComponent,
+    supportedTypes: ['test'],
+    dashboardProps: props,
+  });
+  return null;
+}
+
+it('saves state with usePersistentState hook', async () => {
+  const user = userEvent.setup();
+  const onLayoutConfigChange = jest.fn();
+  let gl: LayoutManager | null = null;
+  const onGoldenLayoutChange = (newGl: LayoutManager) => {
+    gl = newGl;
+  };
+  const { unmount } = makeDashboard({
+    onGoldenLayoutChange,
+    onLayoutConfigChange,
+    children: <TestPlugin />,
+  });
+
+  act(() =>
+    gl!.eventHub.emit(PanelEvent.OPEN, {
+      widget: { type: 'test' },
+    })
+  );
+  expect(screen.getByText('initial')).toBeInTheDocument();
+
+  const waitRAF = () =>
+    act(
+      () =>
+        new Promise(resolve => {
+          requestAnimationFrame(resolve);
+        })
+    );
+
+  // Golden Layout throttles stateChanged events to once per RAF
+  // The test is running too fast and getting batched into one RAF
+  // so we need to wait for the initial stateChanged to fire
+  await waitRAF();
+
+  onLayoutConfigChange.mockClear();
+  await user.click(screen.getByText('initial'));
+  await waitFor(async () =>
+    expect(await screen.findByText('updated')).toBeInTheDocument()
+  );
+
+  // Need to wait here as well because the stateChanged event fires on the next frame
+  // which seems to happen after the unmount without this wait
+  await waitRAF();
+
+  act(() => unmount());
+
+  expect(onLayoutConfigChange).toHaveBeenCalledTimes(1);
+
+  // Remount and verify state is restored
+  makeDashboard({
+    layoutConfig: onLayoutConfigChange.mock.calls[0][0],
+    onLayoutConfigChange,
+    children: <TestPlugin />,
+  });
+
+  await waitFor(() => expect(screen.getByText('updated')).toBeInTheDocument());
+});

packages/dashboard/src/Dashboard.tsx

@@ -47,7 +47,7 @@ export type DashboardProps = {
   dehydrate?: PanelDehydrateFunction;
 
   /** Component to wrap each panel with */
-  panelWrapper?: ComponentType;
+  panelWrapper?: ComponentType<React.PropsWithChildren<PanelProps>>;
 };
 
 export function Dashboard({

packages/dashboard/src/DashboardLayout.tsx

@@ -73,7 +73,7 @@ type DashboardLayoutProps = React.PropsWithChildren<{
   emptyDashboard?: React.ReactNode;
 
   /** Component to wrap each panel with */
-  panelWrapper?: ComponentType<React.PropsWithChildren<DehydratedPanelProps>>;
+  panelWrapper?: ComponentType<React.PropsWithChildren<PanelProps>>;
 }>;
 
 /**
@@ -89,7 +89,7 @@ export function DashboardLayout({
   onLayoutInitialized = DEFAULT_CALLBACK,
   hydrate = hydrateDefault,
   dehydrate = dehydrateDefault,
-  panelWrapper = DashboardPanelWrapper,
+  panelWrapper,
 }: DashboardLayoutProps): JSX.Element {
   const dispatch = useDispatch();
   const data =
@@ -141,22 +141,29 @@ export function DashboardLayout({
          */
         const hasRef = canHaveRef(CType);
 
+        const innerElem = hasRef ? (
+          // eslint-disable-next-line react/jsx-props-no-spreading
+          <CType {...props} ref={ref} />
+        ) : (
+          // eslint-disable-next-line react/jsx-props-no-spreading
+          <CType {...props} />
+        );
+
         // Props supplied by GoldenLayout
         const { glContainer, glEventHub } = props;
         const panelId = LayoutUtils.getIdFromContainer(glContainer);
         return (
           <PanelErrorBoundary glContainer={glContainer} glEventHub={glEventHub}>
             <PanelIdContext.Provider value={panelId as string | null}>
               {/* eslint-disable-next-line react/jsx-props-no-spreading */}
-              <PanelWrapperType {...props}>
-                {hasRef ? (
-                  // eslint-disable-next-line react/jsx-props-no-spreading
-                  <CType {...props} ref={ref} />
+              <DashboardPanelWrapper {...props}>
+                {PanelWrapperType == null ? (
+                  innerElem
                 ) : (
                   // eslint-disable-next-line react/jsx-props-no-spreading
-                  <CType {...props} />
+                  <PanelWrapperType {...props}>{innerElem}</PanelWrapperType>
                 )}
-              </PanelWrapperType>
+              </DashboardPanelWrapper>
             </PanelIdContext.Provider>
           </PanelErrorBoundary>
         );

packages/dashboard/src/DashboardPanelWrapper.tsx

@@ -1,10 +1,42 @@
-import React, { type PropsWithChildren } from 'react';
+import { useCallback, useState, type PropsWithChildren } from 'react';
+import Log from '@deephaven/log';
+import { PersistentStateProvider } from './PersistentStateContext';
+import type { PanelProps } from './DashboardPlugin';
+
+const log = Log.module('DashboardPanelWrapper');
 
 export function DashboardPanelWrapper({
+  glContainer,
   children,
-}: PropsWithChildren<object>): JSX.Element {
-  // eslint-disable-next-line react/jsx-no-useless-fragment
-  return <>{children}</>;
+}: PropsWithChildren<PanelProps>): JSX.Element {
+  const handleDataChange = useCallback(
+    (data: unknown) => {
+      glContainer.setPersistedState(data);
+    },
+    [glContainer]
+  );
+
+  const { persistedState } = glContainer.getConfig();
+
+  // Use a state initializer so we can warn once if the persisted state is invalid
+  const [initialPersistedState] = useState(() => {
+    if (persistedState != null && !Array.isArray(persistedState)) {
+      log.warn(
+        `Persisted state is type ${typeof persistedState}. Expected array. Setting to empty array.`
+      );
+      return [];
+    }
+    return persistedState ?? [];
+  });
+
+  return (
+    <PersistentStateProvider
+      initialState={initialPersistedState}
+      onChange={handleDataChange}
+    >
+      {children}
+    </PersistentStateProvider>
+  );
 }
 
 export default DashboardPanelWrapper;

packages/golden-layout/src/container/ItemContainer.ts

@@ -26,7 +26,10 @@ export default class ItemContainer<
   tab?: Tab;
 
   // This type is to make TS happy and allow ReactComponentConfig passed to container generic
-  _config: C & { componentState: Record<string, unknown> };
+  _config: C & {
+    componentState: Record<string, unknown>;
+    persistedState?: unknown;
+  };
 
   isHidden = false;
 
@@ -192,21 +195,23 @@ export default class ItemContainer<
   }
 
   /**
-   * Merges the provided state into the current one
+   * Notifies the layout manager of a stateupdate
    *
    * @param state
    */
-  extendState(state: string) {
-    this.setState($.extend(true, this.getState(), state));
+  setState(state: Record<string, unknown>) {
+    this._config.componentState = state;
+    this.parent.emitBubblingEvent('stateChanged');
   }
 
   /**
-   * Notifies the layout manager of a stateupdate
+   * Sets persisted state for functional components
+   * which do not have lifecycle methods to hook into.
    *
-   * @param state
+   * @param state The state to persist
    */
-  setState(state: Record<string, unknown>) {
-    this._config.componentState = state;
+  setPersistedState(state: unknown) {
+    this._config.persistedState = state;
     this.parent.emitBubblingEvent('stateChanged');
   }
 

packages/golden-layout/src/utils/ReactComponentHandler.tsx

@@ -1,4 +1,4 @@
-import React from 'react';
+import React, { Component } from 'react';
 import ReactDOM from 'react-dom';
 import $ from 'jquery';
 import type ItemContainer from '../container/ItemContainer';
@@ -21,16 +21,22 @@ export type GLPanelProps = {
 export default class ReactComponentHandler {
   private _container: ItemContainer<ReactComponentConfig>;
 
-  private _reactComponent: React.Component | null = null;
+  private _reactComponent: React.ReactInstance | null = null;
   private _portalComponent: React.ReactPortal | null = null;
-  private _originalComponentWillUpdate: Function | null = null;
+  private _originalComponentWillUpdate:
+    | ((
+        nextProps: Readonly<any>,
+        nextState: Readonly<{}>,
+        nextContext: any
+      ) => void)
+    | undefined = undefined;
   private _initialState: unknown;
   private _reactClass: React.ComponentClass;
 
   constructor(container: ItemContainer<ReactComponentConfig>, state?: unknown) {
     this._reactComponent = null;
     this._portalComponent = null;
-    this._originalComponentWillUpdate = null;
+    this._originalComponentWillUpdate = undefined;
     this._container = container;
     this._initialState = state;
     this._reactClass = this._getReactClass();
@@ -87,18 +93,22 @@ export default class ReactComponentHandler {
    *
    * @param component The component instance created by the `ReactDOM.render` call in the `_render` method.
    */
-  _gotReactComponent(component: React.Component) {
+  _gotReactComponent(component: React.ReactInstance) {
     if (!component) {
       return;
     }
 
     this._reactComponent = component;
-    this._originalComponentWillUpdate =
-      this._reactComponent.componentWillUpdate || function () {};
-    this._reactComponent.componentWillUpdate = this._onUpdate.bind(this);
-    const state = this._container.getState();
-    if (state) {
-      this._reactComponent.setState(state);
+    // Class components manipulate the lifecycle to hook into state changes
+    // Functional components can save data with the usePersistentState hook
+    if (this._reactComponent instanceof Component) {
+      this._originalComponentWillUpdate =
+        this._reactComponent.componentWillUpdate;
+      this._reactComponent.componentWillUpdate = this._onUpdate.bind(this);
+      const state = this._container.getState();
+      if (state) {
+        this._reactComponent.setState(state);
+      }
     }
   }
 
@@ -118,12 +128,13 @@ export default class ReactComponentHandler {
    * Hooks into React's state management and applies the componentstate
    * to GoldenLayout
    */
-  _onUpdate(nextProps: unknown, nextState: Record<string, unknown>) {
+  _onUpdate(nextProps: Readonly<unknown>, nextState: Record<string, unknown>) {
     this._container.setState(nextState);
     this._originalComponentWillUpdate?.call(
       this._reactComponent,
       nextProps,
-      nextState
+      nextState,
+      undefined
     );
   }