Repository: deephaven/deephaven-plugins (language: python)

Issue:
Storing which dashboards/layout/panels are open from this document
Store the placement of the panels/layout that are defined, and upon refresh, re-open the widget into the same panels.


Discussion / hints:

Fix patch:
diff --git a/jest.setup.ts b/jest.setup.ts
index d30aeb930..a4c688a80 100644
--- a/jest.setup.ts
+++ b/jest.setup.ts
@@ -24,3 +24,10 @@ Object.defineProperty(window, 'matchMedia', {
     dispatchEvent: jest.fn(),
   })),
 });
+
+/**
+ * Mock the structuredClone function to use `JSON.stringify` and `JSON.parse`
+ * This is necessary because jsdom does not support `structuredClone`.
+ * https://github.com/jsdom/jsdom/issues/3363
+ */
+global.structuredClone = jest.fn(val => JSON.parse(JSON.stringify(val)));
diff --git a/plugins/ui/src/js/__mocks__/@deephaven/dashboard.js b/plugins/ui/src/js/__mocks__/@deephaven/dashboard.js
index 31e53b69b..bed03cf0c 100644
--- a/plugins/ui/src/js/__mocks__/@deephaven/dashboard.js
+++ b/plugins/ui/src/js/__mocks__/@deephaven/dashboard.js
@@ -1,7 +1,10 @@
 // Mock LayoutUtils, useListener, and PanelEvent from @deephaven/dashboard package
 const mockLayout = {
   root: { contentItems: [], addChild: jest.fn() },
-  eventHub: {},
+  eventHub: {
+    on: jest.fn(),
+    off: jest.fn(),
+  },
   createContentItem: jest.fn(() => ({ setSize: jest.fn() })),
 };
 
@@ -10,6 +13,8 @@ module.exports = {
   ...DashboardActual,
   LayoutUtils: {
     getComponentName: jest.fn(),
+    getStackForConfig: jest.fn(),
+    getIdFromContainer: DashboardActual.LayoutUtils.getIdFromContainer,
     openComponent: jest.fn(),
     closeComponent: jest.fn(),
   },
diff --git a/plugins/ui/src/js/src/DashboardPlugin.tsx b/plugins/ui/src/js/src/DashboardPlugin.tsx
index c3040c7df..9c8373fb6 100644
--- a/plugins/ui/src/js/src/DashboardPlugin.tsx
+++ b/plugins/ui/src/js/src/DashboardPlugin.tsx
@@ -3,6 +3,7 @@ import shortid from 'shortid';
 import {
   DashboardPluginComponentProps,
   LayoutManagerContext,
+  LayoutUtils,
   PanelEvent,
   useListener,
   useDashboardPluginData,
@@ -10,6 +11,7 @@ import {
   WidgetDescriptor,
   PanelOpenEventDetail,
   DEFAULT_DASHBOARD_ID,
+  useDashboardPanel,
 } from '@deephaven/dashboard';
 import Log from '@deephaven/log';
 import {
@@ -18,13 +20,20 @@ import {
 } from '@deephaven/jsapi-bootstrap';
 import { Widget } from '@deephaven/jsapi-types';
 import { ErrorBoundary } from '@deephaven/components';
+import { useDebouncedCallback } from '@deephaven/react-hooks';
 import styles from './styles.scss?inline';
-import { WidgetWrapper } from './widget/WidgetTypes';
+import {
+  ReadonlyWidgetData,
+  WidgetFetch,
+  WidgetId,
+} from './widget/WidgetTypes';
 import PortalPanel from './layout/PortalPanel';
-import WidgetHandler from './widget/WidgetHandler';
+import PortalPanelManager from './layout/PortalPanelManager';
+import DashboardWidgetHandler from './widget/DashboardWidgetHandler';
 
 const NAME_ELEMENT = 'deephaven.ui.Element';
 const DASHBOARD_ELEMENT = 'deephaven.ui.Dashboard';
+const PLUGIN_NAME = '@deephaven/js-plugin-ui.DashboardPlugin';
 
 const log = Log.module('@deephaven/js-plugin-ui.DashboardPlugin');
 
@@ -32,27 +41,45 @@ const log = Log.module('@deephaven/js-plugin-ui.DashboardPlugin');
  * The data stored in redux when the user creates a ui.dashboard.
  */
 interface DashboardPluginData {
-  type: string;
-  title: string;
-  id: string;
+  /** Map of open widgets, along with any data that is stored with them. */
+  openWidgets?: Record<
+    WidgetId,
+    {
+      descriptor: WidgetDescriptor;
+      data?: ReadonlyWidgetData;
+    }
+  >;
+}
+
+interface WidgetWrapper {
+  /** Function to fetch the widget instance from the server */
+  fetch: WidgetFetch;
+
+  /** ID of this widget */
+  id: WidgetId;
+
+  /** Descriptor for the widget. */
   widget: WidgetDescriptor;
+
+  /** Data for the widget */
+  data?: ReadonlyWidgetData;
 }
 
-export function DashboardPlugin({
-  id,
-  layout,
-  registerComponent,
-}: DashboardPluginComponentProps): JSX.Element | null {
-  const [pluginData] = useDashboardPluginData(
+export function DashboardPlugin(
+  props: DashboardPluginComponentProps
+): JSX.Element | null {
+  const { id, layout } = props;
+  const [pluginData, setPluginData] = useDashboardPluginData(
     id,
-    DASHBOARD_ELEMENT
-  ) as unknown as [DashboardPluginData];
+    PLUGIN_NAME
+  ) as unknown as [DashboardPluginData, (data: DashboardPluginData) => void];
+  const [initialPluginData] = useState(pluginData);
 
   const objectFetcher = useObjectFetcher();
 
   // Keep track of the widgets we've got opened.
   const [widgetMap, setWidgetMap] = useState<
-    ReadonlyMap<string, WidgetWrapper>
+    ReadonlyMap<WidgetId, WidgetWrapper>
   >(new Map());
 
   const handleWidgetOpen = useCallback(
@@ -67,7 +94,7 @@ export function DashboardPlugin({
     }) => {
       log.info('Opening widget with ID', widgetId, widget);
       setWidgetMap(prevWidgetMap => {
-        const newWidgetMap = new Map<string, WidgetWrapper>(prevWidgetMap);
+        const newWidgetMap = new Map(prevWidgetMap);
         newWidgetMap.set(widgetId, {
           fetch,
           id: widgetId,
@@ -80,22 +107,19 @@ export function DashboardPlugin({
   );
 
   const handleDashboardOpen = useCallback(
-    ({ widget }: { widget: WidgetDescriptor }) => {
-      const { id: dashboardId, type, name: title = 'Untitled' } = widget;
-      if (dashboardId == null) {
-        log.error("Can't open dashboard without an ID", widget);
-        return;
-      }
-      log.debug('Emitting create dashboard event for', widget);
+    ({
+      widget,
+      dashboardId,
+    }: {
+      widget: WidgetDescriptor;
+      dashboardId: string;
+    }) => {
+      const { name: title = 'Untitled' } = widget;
+      log.debug('Emitting create dashboard event for', dashboardId, widget);
       emitCreateDashboard(layout.eventHub, {
-        pluginId: DASHBOARD_ELEMENT,
+        pluginId: PLUGIN_NAME,
         title,
-        data: {
-          type,
-          title,
-          id: dashboardId,
-          widget,
-        } satisfies DashboardPluginData,
+        data: { openWidgets: { [dashboardId]: { descriptor: widget } } },
       });
     },
     [layout.eventHub]
@@ -116,7 +140,7 @@ export function DashboardPlugin({
           break;
         }
         case DASHBOARD_ELEMENT: {
-          handleDashboardOpen({ widget });
+          handleDashboardOpen({ widget, dashboardId: widgetId });
           break;
         }
         default: {
@@ -128,83 +152,157 @@ export function DashboardPlugin({
   );
 
   useEffect(
-    function loadDashboard() {
-      if (pluginData == null) {
+    function loadInitialPluginData() {
+      if (initialPluginData == null) {
+        log.debug('loadInitialPluginData no data');
         return;
       }
 
-      log.info('Loading dashboard', pluginData);
+      log.debug('loadInitialPluginData', initialPluginData);
 
       setWidgetMap(prevWidgetMap => {
-        const newWidgetMap = new Map<string, WidgetWrapper>(prevWidgetMap);
-        // We need to create a new definition object, otherwise the layout will think it's already open
-        // Can't use a spread operator because the widget definition uses property accessors
-
-        const { widget } = pluginData;
-        newWidgetMap.set(id, {
-          fetch: () => objectFetcher(widget),
-          id,
-          widget,
-        });
+        const newWidgetMap = new Map(prevWidgetMap);
+        const { openWidgets } = initialPluginData;
+        if (openWidgets != null) {
+          Object.entries(openWidgets).forEach(
+            ([widgetId, { descriptor, data }]) => {
+              newWidgetMap.set(widgetId, {
+                fetch: () => objectFetcher(descriptor),
+                id: widgetId,
+                widget: descriptor,
+                data,
+              });
+            }
+          );
+        }
         return newWidgetMap;
       });
     },
-    [objectFetcher, pluginData, id]
+    [objectFetcher, initialPluginData, id]
   );
 
-  const handlePanelClose = useCallback((panelId: string) => {
-    setWidgetMap(prevWidgetMap => {
-      if (!prevWidgetMap.has(panelId)) {
-        return prevWidgetMap;
+  const handlePanelClose = useCallback(
+    (panelId: string) => {
+      log.debug2('handlePanelClose', panelId);
+      setWidgetMap(prevWidgetMap => {
+        if (!prevWidgetMap.has(panelId)) {
+          return prevWidgetMap;
+        }
+        const newWidgetMap = new Map(prevWidgetMap);
+        newWidgetMap.delete(panelId);
+        return newWidgetMap;
+      });
+      // We may need to clean up some panels for this widget if it hasn't actually loaded yet
+      // We should be able to always be able to do this even if it does load, so just remove any panels from the initial load
+      const { openWidgets } = initialPluginData;
+      const openWidget = openWidgets?.[panelId];
+      if (openWidget?.data?.panelIds != null) {
+        const { panelIds } = openWidget.data;
+        for (let i = 0; i < panelIds.length; i += 1) {
+          LayoutUtils.closeComponent(layout.root, { id: panelIds[i] });
+        }
       }
-      const newWidgetMap = new Map<string, WidgetWrapper>(prevWidgetMap);
-      newWidgetMap.delete(panelId);
-      return newWidgetMap;
-    });
-  }, []);
+    },
+    [initialPluginData, layout]
+  );
 
   const handleWidgetClose = useCallback((widgetId: string) => {
-    log.debug('Closing widget', widgetId);
+    log.debug('handleWidgetClose', widgetId);
     setWidgetMap(prevWidgetMap => {
-      const newWidgetMap = new Map<string, WidgetWrapper>(prevWidgetMap);
+      const newWidgetMap = new Map(prevWidgetMap);
       newWidgetMap.delete(widgetId);
       return newWidgetMap;
     });
   }, []);
 
-  useEffect(() => {
-    const cleanups = [registerComponent(PortalPanel.displayName, PortalPanel)];
+  useDashboardPanel({
+    dashboardProps: props,
+    componentName: PortalPanel.displayName,
+    component: PortalPanel,
 
-    return () => {
-      cleanups.forEach(cleanup => cleanup());
-    };
-  }, [registerComponent]);
+    // We don't want these panels to be triggered by a widget opening, we want to control how it is opened later
+    supportedTypes: [],
+  });
 
   // TODO: We need to change up the event system for how objects are opened, since in this case it could be opening multiple panels
   useListener(layout.eventHub, PanelEvent.OPEN, handlePanelOpen);
   useListener(layout.eventHub, PanelEvent.CLOSE, handlePanelClose);
 
+  const sendPluginDataUpdate = useCallback(
+    (newPluginData: DashboardPluginData) => {
+      log.debug('sendPluginDataUpdate', newPluginData);
+      setPluginData(newPluginData);
+    },
+    [setPluginData]
+  );
+
+  const debouncedSendPluginDataUpdate = useDebouncedCallback(
+    sendPluginDataUpdate,
+    500
+  );
+
+  useEffect(
+    function updatePluginData() {
+      // Updates the plugin data with the widgets that are now open in this dashboard
+      const openWidgets: DashboardPluginData['openWidgets'] = {};
+      widgetMap.forEach((widgetWrapper, widgetId) => {
+        openWidgets[widgetId] = {
+          descriptor: widgetWrapper.widget,
+          data: widgetWrapper.data,
+        };
+      });
+      const newPluginData = { openWidgets };
+      debouncedSendPluginDataUpdate(newPluginData);
+    },
+    [widgetMap, debouncedSendPluginDataUpdate]
+  );
+
+  const handleWidgetDataChange = useCallback(
+    (widgetId: string, data: ReadonlyWidgetData) => {
+      log.debug('handleWidgetDataChange', widgetId, data);
+      setWidgetMap(prevWidgetMap => {
+        const newWidgetMap = new Map(prevWidgetMap);
+        const oldWidget = newWidgetMap.get(widgetId);
+        if (oldWidget == null) {
+          throw new Error(`Widget not found: ${widgetId}`);
+        }
+        newWidgetMap.set(widgetId, {
+          ...oldWidget,
+          data,
+        });
+        return newWidgetMap;
+      });
+    },
+    []
+  );
+
   const widgetHandlers = useMemo(
     () =>
-      [...widgetMap.entries()].map(([widgetId, widget]) => (
+      [...widgetMap.entries()].map(([widgetId, wrapper]) => (
         // Fallback to an empty array in default dashboard so we don't display errors over code studio
         <ErrorBoundary
           key={widgetId}
           fallback={id === DEFAULT_DASHBOARD_ID ? [] : null}
         >
-          <DeferredApiBootstrap widget={widget.widget}>
-            <WidgetHandler widget={widget} onClose={handleWidgetClose} />
+          <DeferredApiBootstrap widget={wrapper.widget}>
+            <DashboardWidgetHandler
+              widget={wrapper.widget}
+              id={wrapper.id}
+              initialData={wrapper.data}
+              fetch={wrapper.fetch}
+              onDataChange={handleWidgetDataChange}
+              onClose={handleWidgetClose}
+            />
           </DeferredApiBootstrap>
         </ErrorBoundary>
       )),
-    [handleWidgetClose, widgetMap, id]
+    [handleWidgetClose, handleWidgetDataChange, widgetMap, id]
   );
 
   return (
-    // We'll need to change up how the layout is provided once we have widgets that can open other dashboards...
     <LayoutManagerContext.Provider value={layout}>
       <style>{styles}</style>
-      {widgetHandlers}
+      <PortalPanelManager>{widgetHandlers}</PortalPanelManager>
     </LayoutManagerContext.Provider>
   );
 }
diff --git a/plugins/ui/src/js/src/layout/PortalPanel.tsx b/plugins/ui/src/js/src/layout/PortalPanel.tsx
index 6556720c1..642d7b090 100644
--- a/plugins/ui/src/js/src/layout/PortalPanel.tsx
+++ b/plugins/ui/src/js/src/layout/PortalPanel.tsx
@@ -1,14 +1,7 @@
 import React, { useEffect, useRef } from 'react';
 import { DashboardPanelProps } from '@deephaven/dashboard';
 import { Panel } from '@deephaven/dashboard-core-plugins';
-
-export interface PortalPanelProps extends DashboardPanelProps {
-  /** Listener for when the portal panel is unmounted/closed */
-  onClose: () => void;
-
-  /** Listener for when the portal panel is opened and ready */
-  onOpen: (element: HTMLElement) => void;
-}
+import { emitPortalClosed, emitPortalOpened } from './PortalPanelEvent';
 
 /**
  * Adds and tracks a panel to the GoldenLayout.
@@ -17,9 +10,7 @@ export interface PortalPanelProps extends DashboardPanelProps {
 function PortalPanel({
   glContainer,
   glEventHub,
-  onClose,
-  onOpen,
-}: PortalPanelProps): JSX.Element {
+}: DashboardPanelProps): JSX.Element {
   const ref = useRef<HTMLDivElement>(null);
 
   useEffect(() => {
@@ -27,12 +18,12 @@ function PortalPanel({
     if (current == null) {
       return;
     }
-    onOpen(current);
+    emitPortalOpened(glEventHub, { container: glContainer, element: current });
 
     return () => {
-      onClose();
+      emitPortalClosed(glEventHub, { container: glContainer });
     };
-  }, [onClose, onOpen]);
+  }, [glContainer, glEventHub]);
 
   return (
     <Panel glContainer={glContainer} glEventHub={glEventHub}>
diff --git a/plugins/ui/src/js/src/layout/PortalPanelEvent.ts b/plugins/ui/src/js/src/layout/PortalPanelEvent.ts
new file mode 100644
index 000000000..d187c76fd
--- /dev/null
+++ b/plugins/ui/src/js/src/layout/PortalPanelEvent.ts
@@ -0,0 +1,119 @@
+import { DashboardPanelProps } from '@deephaven/dashboard';
+import { useEffect } from 'react';
+
+/**
+ * Emitted when a portal panel is opened
+ */
+export const PORTAL_OPENED = 'PortalPanelEvent.PORTAL_OPENED';
+
+/**
+ * Emitted when a portal panel is closed
+ */
+export const PORTAL_CLOSED = 'PortalPanelEvent.PORTAL_CLOSED';
+
+export type EventListenerRemover = () => void;
+export type EventListenFunction<TPayload = unknown> = (
+  eventHub: DashboardPanelProps['glEventHub'],
+  handler: (p: TPayload) => void
+) => EventListenerRemover;
+
+export type EventEmitFunction<TPayload = unknown> = (
+  eventHub: DashboardPanelProps['glEventHub'],
+  payload: TPayload
+) => void;
+
+export type EventListenerHook<TPayload = unknown> = (
+  eventHub: DashboardPanelProps['glEventHub'],
+  handler: (p: TPayload) => void
+) => void;
+
+/**
+ * Listen for an event
+ * @param eventHub The event hub to listen to
+ * @param event The event to listen for
+ * @param handler The handler to call when the event is emitted
+ * @returns A function to stop listening for the event
+ */
+export function listenForEvent<TPayload>(
+  eventHub: DashboardPanelProps['glEventHub'],
+  event: string,
+  handler: (p: TPayload) => void
+): EventListenerRemover {
+  eventHub.on(event, handler);
+  return () => {
+    eventHub.off(event, handler);
+  };
+}
+
+export function makeListenFunction<TPayload>(
+  event: string
+): EventListenFunction<TPayload> {
+  return (eventHub, handler) => listenForEvent(eventHub, event, handler);
+}
+
+export function makeEmitFunction<TPayload>(
+  event: string
+): EventEmitFunction<TPayload> {
+  return (eventHub, payload) => {
+    eventHub.emit(event, payload);
+  };
+}
+
+export function makeUseListenerFunction<TPayload>(
+  event: string
+): EventListenerHook<TPayload> {
+  return (eventHub, handler) => {
+    useEffect(
+      () => listenForEvent(eventHub, event, handler),
+      [eventHub, handler]
+    );
+  };
+}
+
+/**
+ * Create listener, emitter, and hook functions for an event
+ * @param event Name of the event to create functions for
+ * @returns Listener, Emitter, and Hook functions for the event
+ */
+export function makeEventFunctions<TPayload>(event: string): {
+  listen: EventListenFunction<TPayload>;
+  emit: EventEmitFunction<TPayload>;
+  useListener: EventListenerHook<TPayload>;
+} {
+  return {
+    listen: makeListenFunction<TPayload>(event),
+    emit: makeEmitFunction<TPayload>(event),
+    useListener: makeUseListenerFunction<TPayload>(event),
+  };
+}
+
+export interface PortalOpenedPayload {
+  /**
+   * Golden Layout Container object.
+   * Can get the ID of the container using `LayoutUtils.getIdFromContainer` to identify this container.
+   */
+  container: DashboardPanelProps['glContainer'];
+
+  /** Element to use for the portal */
+  element: HTMLElement;
+}
+
+export interface PortalClosedPayload {
+  /**
+   * Golden Layout Container object.
+   * Can get the ID of the container using `LayoutUtils.getIdFromContainer` to identify this container.
+   */
+  container: DashboardPanelProps['glContainer'];
+}
+
+export const {
+  listen: listenForPortalOpened,
+  emit: emitPortalOpened,
+  useListener: usePortalOpenedListener,
+} = makeEventFunctions<PortalOpenedPayload>(PORTAL_OPENED);
+
+export const {
+  listen: listenForPortalClosed,
+  emit: emitPortalClosed,
+  useListener: usePortalClosedListener,
+} = makeEventFunctions<PortalClosedPayload>(PORTAL_CLOSED);
diff --git a/plugins/ui/src/js/src/layout/PortalPanelManager.tsx b/plugins/ui/src/js/src/layout/PortalPanelManager.tsx
new file mode 100644
index 000000000..4533ac9d6
--- /dev/null
+++ b/plugins/ui/src/js/src/layout/PortalPanelManager.tsx
@@ -0,0 +1,60 @@
+import React, { useCallback, useState } from 'react';
+import { LayoutUtils, useLayoutManager } from '@deephaven/dashboard';
+import {
+  PortalPanelMap,
+  PortalPanelManagerContext,
+} from './PortalPanelManagerContext';
+import {
+  usePortalClosedListener,
+  usePortalOpenedListener,
+} from './PortalPanelEvent';
+
+/**
+ * Listens for PortalPanels being opened and closed, and maintains a map of currently open portal elements.
+ * Sets this in the PortalPanelManagerContext for downstream consumption.
+ */
+function PortalPanelManager({
+  children,
+}: React.PropsWithChildren<unknown>): JSX.Element {
+  const { eventHub } = useLayoutManager();
+  const [portals, setPortals] = useState<PortalPanelMap>(new Map());
+
+  const handlePortalOpened = useCallback(({ container, element }) => {
+    setPortals(prevPortals => {
+      const containerId = LayoutUtils.getIdFromContainer(container);
+      if (containerId == null) {
+        throw new Error('Invalid panel ID');
+      }
+
+      const panelId = Array.isArray(containerId) ? containerId[0] : containerId;
+      const newPortals = new Map(prevPortals);
+      newPortals.set(panelId, element);
+      return newPortals;
+    });
+  }, []);
+
+  const handlePortalClosed = useCallback(({ container }) => {
+    setPortals(prevPortals => {
+      const containerId = LayoutUtils.getIdFromContainer(container);
+      if (containerId == null) {
+        throw new Error('Invalid panel ID');
+      }
+
+      const panelId = Array.isArray(containerId) ? containerId[0] : containerId;
+      const newPortals = new Map(prevPortals);
+      newPortals.delete(panelId);
+      return newPortals;
+    });
+  }, []);
+
+  usePortalOpenedListener(eventHub, handlePortalOpened);
+  usePortalClosedListener(eventHub, handlePortalClosed);
+
+  return (
+    <PortalPanelManagerContext.Provider value={portals}>
+      {children}
+    </PortalPanelManagerContext.Provider>
+  );
+}
+
+export default PortalPanelManager;
diff --git a/plugins/ui/src/js/src/layout/PortalPanelManagerContext.ts b/plugins/ui/src/js/src/layout/PortalPanelManagerContext.ts
new file mode 100644
index 000000000..7ecc58902
--- /dev/null
+++ b/plugins/ui/src/js/src/layout/PortalPanelManagerContext.ts
@@ -0,0 +1,14 @@
+import React from 'react';
+import { useContextOrThrow } from '@deephaven/react-hooks';
+
+/** Map from the panel IDs to the element for that panel */
+export type PortalPanelMap = ReadonlyMap<string, HTMLElement>;
+
+export const PortalPanelManagerContext =
+  React.createContext<PortalPanelMap | null>(null);
+
+export function usePortalPanelManager() {
+  return useContextOrThrow(PortalPanelManagerContext, 'PortalPanelManager');
+}
+
+export default PortalPanelManagerContext;
diff --git a/plugins/ui/src/js/src/layout/ReactPanel.tsx b/plugins/ui/src/js/src/layout/ReactPanel.tsx
index a96f4f771..d77037a8b 100644
--- a/plugins/ui/src/js/src/layout/ReactPanel.tsx
+++ b/plugins/ui/src/js/src/layout/ReactPanel.tsx
@@ -1,12 +1,5 @@
-import React, {
-  useCallback,
-  useEffect,
-  useMemo,
-  useRef,
-  useState,
-} from 'react';
+import React, { useCallback, useEffect, useRef } from 'react';
 import ReactDOM from 'react-dom';
-import shortid from 'shortid';
 import {
   LayoutUtils,
   PanelEvent,
@@ -15,10 +8,11 @@ import {
 } from '@deephaven/dashboard';
 import Log from '@deephaven/log';
 import PortalPanel from './PortalPanel';
-import { useReactPanelManager } from './ReactPanelManager';
+import { ReactPanelControl, useReactPanel } from './ReactPanelManager';
 import { ReactPanelProps } from './LayoutUtils';
 import { useParentItem } from './ParentItemContext';
 import { ReactPanelContext } from './ReactPanelContext';
+import { usePortalPanelManager } from './PortalPanelManagerContext';
 
 const log = Log.module('@deephaven/js-plugin-ui/ReactPanel');
 
@@ -27,13 +21,18 @@ const log = Log.module('@deephaven/js-plugin-ui/ReactPanel');
  */
 function ReactPanel({ children, title }: ReactPanelProps) {
   const layoutManager = useLayoutManager();
-  const panelManager = useReactPanelManager();
-  const { metadata, onClose, onOpen } = panelManager;
-  const panelId = useMemo(() => shortid(), []);
-  const [element, setElement] = useState<HTMLElement>();
-  const isPanelOpenRef = useRef(false);
-  const openedMetadataRef = useRef<Record<string, unknown>>();
+  const { metadata, onClose, onOpen, panelId } = useReactPanel();
+  const portalManager = usePortalPanelManager();
+  const portal = portalManager.get(panelId);
+
+  // If there is already a portal that exists, then we're rehydrating from a dehydrated state
+  // Initialize the `isPanelOpenRef` and `openedWidgetRef` accordingly on initialization
+  const isPanelOpenRef = useRef(portal != null);
+  const openedMetadataRef = useRef<ReactPanelControl['metadata']>(
+    portal != null ? metadata : null
+  );
   const parent = useParentItem();
+  const { eventHub } = layoutManager;
 
   log.debug2('Rendering panel', panelId);
 
@@ -43,7 +42,7 @@ function ReactPanel({ children, title }: ReactPanelProps) {
         log.debug('Closing panel', panelId);
         LayoutUtils.closeComponent(parent, { id: panelId });
         isPanelOpenRef.current = false;
-        onClose(panelId);
+        onClose();
       }
     },
     [parent, onClose, panelId]
@@ -54,52 +53,58 @@ function ReactPanel({ children, title }: ReactPanelProps) {
       if (closedPanelId === panelId) {
         log.debug('Panel closed', panelId);
         isPanelOpenRef.current = false;
-        onClose(panelId);
+        onClose();
       }
     },
     [onClose, panelId]
   );
 
-  useListener(layoutManager.eventHub, PanelEvent.CLOSED, handlePanelClosed);
+  useListener(eventHub, PanelEvent.CLOSED, handlePanelClosed);
 
-  useEffect(() => {
-    if (
-      isPanelOpenRef.current === false ||
-      openedMetadataRef.current !== metadata
-    ) {
-      const panelTitle =
-        title ?? (typeof metadata?.name === 'string' ? metadata.name : '');
-      const config = {
-        type: 'react-component' as const,
-        component: PortalPanel.displayName,
-        props: {
+  useEffect(
+    /** Opens a panel in the layout if necessary. Triggered when the panel metadata changes or the panel has not been opened yet. */
+    function openIfNecessary() {
+      if (isPanelOpenRef.current === false) {
+        const existingStack = LayoutUtils.getStackForConfig(parent, {
           id: panelId,
-          onClose: () => {
-            isPanelOpenRef.current = false;
-            setElement(undefined);
-          },
-          onOpen: setElement,
-          metadata,
-        },
-        title: panelTitle,
-        id: panelId,
-      };
+        });
+        if (existingStack != null) {
+          log.debug2('Panel already exists, just re-using');
+          isPanelOpenRef.current = true;
+          return;
+        }
+      }
 
-      LayoutUtils.openComponent({ root: parent, config });
-      log.debug('Opened panel', panelId, config);
-      isPanelOpenRef.current = true;
-      openedMetadataRef.current = metadata;
+      if (
+        isPanelOpenRef.current === false ||
+        openedMetadataRef.current !== metadata
+      ) {
+        const panelTitle = title ?? metadata?.name ?? '';
+        const config = {
+          type: 'react-component' as const,
+          component: PortalPanel.displayName,
+          props: {},
+          title: panelTitle,
+          id: panelId,
+        };
 
-      onOpen(panelId);
-    }
-  }, [parent, metadata, onOpen, panelId, title]);
+        LayoutUtils.openComponent({ root: parent, config });
+        log.debug('Opened panel', panelId, config);
+        isPanelOpenRef.current = true;
+        openedMetadataRef.current = metadata;
+
+        onOpen();
+      }
+    },
+    [parent, metadata, onOpen, panelId, title]
+  );
 
-  return element
+  return portal
     ? ReactDOM.createPortal(
         <ReactPanelContext.Provider value={panelId}>
           {children}
         </ReactPanelContext.Provider>,
-        element
+        portal
       )
     : null;
 }
diff --git a/plugins/ui/src/js/src/layout/ReactPanelManager.ts b/plugins/ui/src/js/src/layout/ReactPanelManager.ts
index 53db36f60..2bee80f44 100644
--- a/plugins/ui/src/js/src/layout/ReactPanelManager.ts
+++ b/plugins/ui/src/js/src/layout/ReactPanelManager.ts
@@ -1,26 +1,71 @@
-import { createContext, useContext } from 'react';
+import { PanelProps } from '@deephaven/dashboard';
+import { useContextOrThrow } from '@deephaven/react-hooks';
+import { createContext, useCallback, useMemo } from 'react';
 
-export type ReactPanelManager = {
+/**
+ * Manager for panels within a widget. This is used to manage the lifecycle of panels within a widget.
+ */
+export interface ReactPanelManager {
   /**
-   * Metadata to pass to all the panels.
+   * Metadata stored with the panel. Typically a descriptor of the widget opening the panel and used for hydration.
    * Updating the metadata will cause the panel to be re-opened, or replaced if it is closed.
    * Can also be used for rehydration.
    */
-  metadata: Record<string, unknown>;
+  metadata: PanelProps['metadata'];
 
   /** Triggered when a panel is opened */
   onOpen: (panelId: string) => void;
 
   /** Triggered when a panel is closed */
   onClose: (panelId: string) => void;
-};
 
-export const ReactPanelManagerContext = createContext<ReactPanelManager>({
-  metadata: { name: '', type: '' },
-  onOpen: () => undefined,
-  onClose: () => undefined,
-});
+  /**
+   * Get a unique panelId from the panel manager. This should be used to identify the panel in the layout.
+   */
+  getPanelId: () => string;
+}
+
+/** Interface for using a react panel */
+export interface ReactPanelControl {
+  /**
+   * Metadata stored with the panel. Typically a descriptor of the widget opening the panel and used for hydration.
+   * Updating the metadata will cause the panel to be re-opened, or replaced if it is closed.
+   * Can also be used for rehydration.
+   */
+  metadata: PanelProps['metadata'];
+
+  /** Must be called when the panel is opened */
+  onOpen: () => void;
+
+  /** Must be called when the panel is closed */
+  onClose: () => void;
+
+  /** The panelId for this react panel */
+  panelId: string;
+}
+
+export const ReactPanelManagerContext = createContext<ReactPanelManager | null>(
+  null
+);
 
 export function useReactPanelManager(): ReactPanelManager {
-  return useContext(ReactPanelManagerContext);
+  return useContextOrThrow(
+    ReactPanelManagerContext,
+    'No ReactPanelManager found, did you wrap in a ReactPanelManagerProvider.Context?'
+  );
+}
+
+/**
+ * Use the controls for a single react panel.
+ */
+export function useReactPanel(): ReactPanelControl {
+  const { metadata, onClose, onOpen, getPanelId } = useReactPanelManager();
+  const panelId = useMemo(() => getPanelId(), [getPanelId]);
+
+  return {
+    metadata,
+    onClose: useCallback(() => onClose(panelId), [onClose, panelId]),
+    onOpen: useCallback(() => onOpen(panelId), [onOpen, panelId]),
+    panelId,
+  };
 }
diff --git a/plugins/ui/src/js/src/layout/ReactPanelManagerProvider.tsx b/plugins/ui/src/js/src/layout/ReactPanelManagerProvider.tsx
new file mode 100644
index 000000000..990341f75
--- /dev/null
+++ b/plugins/ui/src/js/src/layout/ReactPanelManagerProvider.tsx
@@ -0,0 +1,30 @@
+import React, { useMemo } from 'react';
+import {
+  ReactPanelManager,
+  ReactPanelManagerContext,
+} from './ReactPanelManager';
+
+export function ReactPanelManagerProvider({
+  children,
+  metadata,
+  onOpen,
+  onClose,
+  getPanelId,
+}: React.PropsWithChildren<ReactPanelManager>): JSX.Element {
+  const manager = useMemo(
+    () => ({
+      metadata,
+      onOpen,
+      onClose,
+      getPanelId,
+    }),
+    [metadata, onOpen, onClose, getPanelId]
+  );
+  return (
+    <ReactPanelManagerContext.Provider value={manager}>
+      {children}
+    </ReactPanelManagerContext.Provider>
+  );
+}
+
+export default ReactPanelManagerProvider;
diff --git a/plugins/ui/src/js/src/widget/DashboardWidgetHandler.tsx b/plugins/ui/src/js/src/widget/DashboardWidgetHandler.tsx
new file mode 100644
index 000000000..26660e569
--- /dev/null
+++ b/plugins/ui/src/js/src/widget/DashboardWidgetHandler.tsx
@@ -0,0 +1,65 @@
+/**
+ * Handles document events for one widget.
+ */
+import React, { useCallback } from 'react';
+import { WidgetDescriptor } from '@deephaven/dashboard';
+import { Widget } from '@deephaven/jsapi-types';
+import Log from '@deephaven/log';
+import { ReadonlyWidgetData, WidgetId } from './WidgetTypes';
+import WidgetHandler from './WidgetHandler';
+
+const log = Log.module('@deephaven/js-plugin-ui/DashboardWidgetHandler');
+
+export interface DashboardWidgetHandlerProps {
+  /** ID of this widget */
+  id: WidgetId;
+
+  /** Widget for this to handle */
+  widget: WidgetDescriptor;
+
+  /** Fetch the widget instance */
+  fetch: () => Promise<Widget>;
+
+  /** Widget data to display */
+  initialData?: ReadonlyWidgetData;
+
+  /** Triggered when all panels opened from this widget have closed */
+  onClose?: (widgetId: WidgetId) => void;
+
+  /** Triggered when the data in the widget changes */
+  onDataChange?: (widgetId: WidgetId, data: ReadonlyWidgetData) => void;
+}
+
+function DashboardWidgetHandler({
+  id,
+  onClose,
+  onDataChange,
+  ...otherProps
+}: DashboardWidgetHandlerProps) {
+  const handleClose = useCallback(() => {
+    log.debug('handleClose', id);
+    onClose?.(id);
+  }, [onClose, id]);
+
+  const handleDataChange = useCallback(
+    (data: ReadonlyWidgetData) => {
+      log.debug('handleDataChange', id, data);
+      onDataChange?.(id, data);
+    },
+    [onDataChange, id]
+  );
+
+  return (
+    <WidgetHandler
+      onDataChange={handleDataChange}
+      onClose={handleClose}
+      // eslint-disable-next-line react/jsx-props-no-spreading
+      {...otherProps}
+    />
+  );
+}
+
+DashboardWidgetHandler.displayName =
+  '@deephaven/js-plugin-ui/DashboardWidgetHandler';
+
+export default DashboardWidgetHandler;
diff --git a/plugins/ui/src/js/src/widget/DocumentHandler.tsx b/plugins/ui/src/js/src/widget/DocumentHandler.tsx
index 0419f5e30..4cd2485a6 100644
--- a/plugins/ui/src/js/src/widget/DocumentHandler.tsx
+++ b/plugins/ui/src/js/src/widget/DocumentHandler.tsx
@@ -1,15 +1,29 @@
-import React, { useCallback, useMemo, useRef } from 'react';
+import React, { useCallback, useMemo, useRef, useState } from 'react';
+import shortid from 'shortid';
 import { WidgetDescriptor } from '@deephaven/dashboard';
 import Log from '@deephaven/log';
+import { EMPTY_FUNCTION } from '@deephaven/utils';
 import { ReactPanelManagerContext } from '../layout/ReactPanelManager';
 import { getRootChildren } from './DocumentUtils';
+import { ReadonlyWidgetData, WidgetData } from './WidgetTypes';
 
 const log = Log.module('@deephaven/js-plugin-ui/DocumentHandler');
 
+const EMPTY_OBJECT = Object.freeze({});
+
 export type DocumentHandlerProps = React.PropsWithChildren<{
   /** Definition of the widget used to create this document. Used for titling panels if necessary. */
   widget: WidgetDescriptor;
 
+  /**
+   * Data state to use when loading the widget.
+   * When the data state is updated, the new state is emitted via the `onDataChange` callback.
+   */
+  initialData?: ReadonlyWidgetData;
+
+  /** Triggered when the data in the document changes */
+  onDataChange?: (data: ReadonlyWidgetData) => void;
+
   /** Triggered when all panels opened from this document have closed */
   onClose?: () => void;
 }>;
@@ -20,41 +34,70 @@ export type DocumentHandlerProps = React.PropsWithChildren<{
  * or all non-panels, which will automatically be wrapped in one panel.
  * Responsible for opening any panels or dashboards specified in the document.
  */
-function DocumentHandler({ children, widget, onClose }: DocumentHandlerProps) {
+function DocumentHandler({
+  children,
+  widget,
+  initialData: data = EMPTY_OBJECT,
+  onDataChange = EMPTY_FUNCTION,
+  onClose,
+}: DocumentHandlerProps) {
   log.debug('Rendering document', widget);
   const panelOpenCountRef = useRef(0);
+  const panelIdIndex = useRef(0);
+  const [widgetData] = useState<WidgetData>(() => structuredClone(data));
 
-  const metadata = useMemo(
-    () => ({
-      name: widget.name ?? 'Unknown',
-      type: widget.type,
-    }),
-    [widget]
+  const handleOpen = useCallback(
+    (panelId: string) => {
+      panelOpenCountRef.current += 1;
+      log.debug('Panel opened, open count', panelOpenCountRef.current);
+
+      if (widgetData.panelIds == null) {
+        widgetData.panelIds = [];
+      }
+      widgetData.panelIds?.push(panelId);
+      onDataChange(widgetData);
+    },
+    [onDataChange, widgetData]
+  );
+
+  const handleClose = useCallback(
+    (panelId: string) => {
+      panelOpenCountRef.current -= 1;
+      if (panelOpenCountRef.current < 0) {
+        throw new Error('Panel open count is negative');
+      }
+      log.debug('Panel closed, open count', panelOpenCountRef.current);
+      if (panelOpenCountRef.current === 0) {
+        onClose?.();
+        return;
+      }
+
+      widgetData.panelIds = (widgetData.panelIds ?? [])?.filter(
+        id => id !== panelId
+      );
+      onDataChange(widgetData);
+    },
+    [onClose, onDataChange, widgetData]
   );
 
-  const handleOpen = useCallback(() => {
-    panelOpenCountRef.current += 1;
-    log.debug('Panel opened, open count', panelOpenCountRef.current);
-  }, []);
-
-  const handleClose = useCallback(() => {
-    panelOpenCountRef.current -= 1;
-    if (panelOpenCountRef.current < 0) {
-      throw new Error('Panel open count is negative');
-    }
-    log.debug('Panel closed, open count', panelOpenCountRef.current);
-    if (panelOpenCountRef.current === 0) {
-      onClose?.();
-    }
-  }, [onClose]);
+  const getPanelId = useCallback(() => {
+    // On rehydration, yield known IDs first
+    // If there are no more known IDs, generate a new one.
+    // This can happen if the document hasn't been opened before, or if it's rehydrated and a new panel is added.
+    // Note that if the order of panels changes, the worst case scenario is that panels appear in the wrong location in the layout.
+    const panelId = widgetData.panelIds?.[panelIdIndex.current] ?? shortid();
+    panelIdIndex.current += 1;
+    return panelId;
+  }, [widgetData]);
 
   const panelManager = useMemo(
     () => ({
-      metadata,
+      metadata: widget,
       onOpen: handleOpen,
       onClose: handleClose,
+      getPanelId,
     }),
-    [metadata, handleClose, handleOpen]
+    [widget, getPanelId, handleClose, handleOpen]
   );
 
   return (
diff --git a/plugins/ui/src/js/src/widget/WidgetHandler.tsx b/plugins/ui/src/js/src/widget/WidgetHandler.tsx
index f7841860d..0e2e1052e 100644
--- a/plugins/ui/src/js/src/widget/WidgetHandler.tsx
+++ b/plugins/ui/src/js/src/widget/WidgetHandler.tsx
@@ -14,8 +14,10 @@ import {
   JSONRPCServer,
   JSONRPCServerAndClient,
 } from 'json-rpc-2.0';
+import { WidgetDescriptor } from '@deephaven/dashboard';
 import type { Widget, WidgetExportedObject } from '@deephaven/jsapi-types';
 import Log from '@deephaven/log';
+import { EMPTY_FUNCTION } from '@deephaven/utils';
 import {
   CALLABLE_KEY,
   OBJECT_KEY,
@@ -23,7 +25,7 @@ import {
   isElementNode,
   isObjectNode,
 } from '../elements/ElementUtils';
-import { WidgetMessageEvent, WidgetWrapper } from './WidgetTypes';
+import { ReadonlyWidgetData, WidgetMessageEvent } from './WidgetTypes';
 import DocumentHandler from './DocumentHandler';
 import { getComponentForElement } from './WidgetUtils';
 
@@ -31,13 +33,28 @@ const log = Log.module('@deephaven/js-plugin-ui/WidgetHandler');
 
 export interface WidgetHandlerProps {
   /** Widget for this to handle */
-  widget: WidgetWrapper;
+  widget: WidgetDescriptor;
+
+  /** Fetch the widget instance */
+  fetch: () => Promise<Widget>;
+
+  /** Widget data to display */
+  initialData?: ReadonlyWidgetData;
 
   /** Triggered when all panels opened from this widget have closed */
-  onClose?: (widgetId: string) => void;
+  onClose?: () => void;
+
+  /** Triggered when the data in the widget changes */
+  onDataChange?: (data: ReadonlyWidgetData) => void;
 }
 
-function WidgetHandler({ onClose, widget: wrapper }: WidgetHandlerProps) {
+function WidgetHandler({
+  onClose,
+  onDataChange = EMPTY_FUNCTION,
+  fetch,
+  widget: descriptor,
+  initialData,
+}: WidgetHandlerProps) {
   const [widget, setWidget] = useState<Widget>();
   const [document, setDocument] = useState<ReactNode>();
 
@@ -213,10 +230,10 @@ function WidgetHandler({ onClose, widget: wrapper }: WidgetHandlerProps) {
 
   useEffect(
     function loadWidget() {
-      log.debug('loadWidget', wrapper.id, wrapper.widget);
+      log.debug('loadWidget', descriptor);
       let isCancelled = false;
       async function loadWidgetInternal() {
-        const newWidget = await wrapper.fetch();
+        const newWidget = await fetch();
         if (isCancelled) {
           newWidget.close();
           newWidget.exportedObjects.forEach(exportedObject => {
@@ -224,7 +241,7 @@ function WidgetHandler({ onClose, widget: wrapper }: WidgetHandlerProps) {
           });
           return;
         }
-        log.debug('newWidget', wrapper.id, wrapper.widget, newWidget);
+        log.debug('newWidget', descriptor, newWidget);
         setWidget(newWidget);
       }
       loadWidgetInternal();
@@ -232,22 +249,22 @@ function WidgetHandler({ onClose, widget: wrapper }: WidgetHandlerProps) {
         isCancelled = true;
       };
     },
-    [wrapper]
+    [fetch, descriptor]
   );
 
-  const handleDocumentClose = useCallback(() => {
-    log.debug('Widget document closed', wrapper.id);
-    onClose?.(wrapper.id);
-  }, [onClose, wrapper.id]);
-
   return useMemo(
     () =>
       document != null ? (
-        <DocumentHandler widget={wrapper.widget} onClose={handleDocumentClose}>
+        <DocumentHandler
+          widget={descriptor}
+          initialData={initialData}
+          onDataChange={onDataChange}
+          onClose={onClose}
+        >
           {document}
         </DocumentHandler>
       ) : null,
-    [document, handleDocumentClose, wrapper]
+    [document, descriptor, initialData, onClose, onDataChange]
   );
 }
 
diff --git a/plugins/ui/src/js/src/widget/WidgetTestUtils.ts b/plugins/ui/src/js/src/widget/WidgetTestUtils.ts
index 60e776620..11afff8b8 100644
--- a/plugins/ui/src/js/src/widget/WidgetTestUtils.ts
+++ b/plugins/ui/src/js/src/widget/WidgetTestUtils.ts
@@ -1,7 +1,6 @@
 import { WidgetDescriptor } from '@deephaven/dashboard';
 import { TestUtils } from '@deephaven/utils';
 import type { Widget } from '@deephaven/jsapi-types';
-import { WidgetWrapper } from './WidgetTypes';
 
 export function makeDocumentUpdatedJsonRpc(
   document: Record<string, unknown> = {}
@@ -42,14 +41,3 @@ export function makeWidget({
     exportedObjects,
   });
 }
-
-export function makeWidgetWrapper({
-  widget = makeWidgetDescriptor(),
-  fetch = () => Promise.resolve(makeWidget()),
-}: Partial<WidgetWrapper> = {}): WidgetWrapper {
-  return {
-    id: widget.id ?? 'widget-id',
-    widget,
-    fetch,
-  };
-}
diff --git a/plugins/ui/src/js/src/widget/WidgetTypes.ts b/plugins/ui/src/js/src/widget/WidgetTypes.ts
index ab935c763..f56548dce 100644
--- a/plugins/ui/src/js/src/widget/WidgetTypes.ts
+++ b/plugins/ui/src/js/src/widget/WidgetTypes.ts
@@ -1,6 +1,7 @@
-import { WidgetDescriptor } from '@deephaven/dashboard';
 import { Widget, WidgetExportedObject } from '@deephaven/jsapi-types';
 
+export type WidgetId = string;
+
 export interface WidgetMessageDetails {
   getDataAsBase64(): string;
   getDataAsString(): string;
@@ -11,8 +12,9 @@ export type WidgetMessageEvent = CustomEvent<WidgetMessageDetails>;
 
 export type WidgetFetch = (takeOwnership?: boolean) => Promise<Widget>;
 
-export type WidgetWrapper = {
-  fetch: WidgetFetch;
-  id: string;
-  widget: WidgetDescriptor;
+export type WidgetData = {
+  /** Panel IDs that are opened by this widget */
+  panelIds?: string[];
 };
+
+export type ReadonlyWidgetData = Readonly<WidgetData>;


Test patch:
diff --git a/plugins/ui/src/js/src/layout/PortalPanelManager.test.tsx b/plugins/ui/src/js/src/layout/PortalPanelManager.test.tsx
new file mode 100644
index 000000000..62e7a9292
--- /dev/null
+++ b/plugins/ui/src/js/src/layout/PortalPanelManager.test.tsx
@@ -0,0 +1,129 @@
+import React from 'react';
+import { act, render } from '@testing-library/react';
+import PortalPanelManager from './PortalPanelManager';
+
+// Mock the usePortalOpenedListener and usePortalClosedListener functions
+const mockUsePortalOpenedListener = jest.fn();
+const mockUsePortalClosedListener = jest.fn();
+jest.mock('./PortalPanelEvent', () => ({
+  usePortalClosedListener: jest.fn((...args) => {
+    mockUsePortalClosedListener(...args);
+  }),
+  usePortalOpenedListener: jest.fn((...args) => {
+    mockUsePortalOpenedListener(...args);
+  }),
+}));
+
+const mockProvider = jest.fn(({ children }) => children);
+jest.mock('./PortalPanelManagerContext', () => ({
+  PortalPanelManagerContext: {
+    Provider: jest.fn(props => mockProvider(props)),
+  },
+}));
+
+beforeEach(() => {
+  jest.clearAllMocks();
+});
+
+describe('PortalPanelManager', () => {
+  it('should call the usePortalOpenedListener and usePortalClosedListener with the correct callbacks', () => {
+    // Render the PortalPanelManager component
+    render(
+      <PortalPanelManager>
+        <div>Test</div>
+      </PortalPanelManager>
+    );
+
+    // Verify that the usePortalOpenedListener and usePortalClosedListener functions are called with the correct callbacks
+    expect(mockUsePortalOpenedListener).toHaveBeenCalledTimes(1);
+    expect(mockUsePortalClosedListener).toHaveBeenCalledTimes(1);
+  });
+
+  it('should render the children wrapped in the PortalPanelManagerContext.Provider', () => {
+    // Render the PortalPanelManager component with children
+    const { getByText } = render(
+      <PortalPanelManager>
+        <div>Test</div>
+      </PortalPanelManager>
+    );
+
+    // Verify that the children are rendered and wrapped in the PortalPanelManagerContext.Provider
+    expect(getByText('Test')).toBeInTheDocument();
+  });
+
+  it('should add portals to the context when they are opened', () => {
+    const mockContainer1 = { _config: { id: 'test-container-1' } };
+    const mockElement1 = document.createElement('div');
+    const mockContainer2 = { _config: { id: 'test-container-2' } };
+    const mockElement2 = document.createElement('div');
+
+    // Render the PortalPanelManager component
+    render(
+      <PortalPanelManager>
+        <div>Test</div>
+      </PortalPanelManager>
+    );
+
+    // Verify that the setPortals function is called when a portal is opened
+    expect(mockProvider).toHaveBeenCalledWith(
+      expect.objectContaining({ value: new Map() })
+    );
+    mockProvider.mockClear();
+
+    act(() => {
+      mockUsePortalOpenedListener.mock.calls[0][1]({
+        container: mockContainer1,
+        element: mockElement1,
+      });
+    });
+
+    expect(mockProvider).toHaveBeenCalledWith(
+      expect.objectContaining({
+        value: new Map([[mockContainer1, mockElement1]]),
+      })
+    );
+    mockProvider.mockClear();
+
+    act(() => {
+      mockUsePortalOpenedListener.mock.calls[0][1]({
+        container: mockContainer2,
+        element: mockElement2,
+      });
+    });
+
+    expect(mockProvider).toHaveBeenCalledWith(
+      expect.objectContaining({
+        value: new Map([
+          [mockContainer1, mockElement1],
+          [mockContainer2, mockElement2],
+        ]),
+      })
+    );
+    mockProvider.mockClear();
+
+    // Verify the mock provider gets updated portals when portals are closed
+    act(() => {
+      mockUsePortalClosedListener.mock.calls[0][1]({
+        container: mockContainer1,
+      });
+    });
+
+    expect(mockProvider).toHaveBeenCalledWith(
+      expect.objectContaining({
+        value: new Map([[mockContainer2, mockElement2]]),
+      })
+    );
+    mockProvider.mockClear();
+
+    // Close the final portal
+    act(() => {
+      mockUsePortalClosedListener.mock.calls[0][1]({
+        container: mockContainer2,
+      });
+    });
+
+    expect(mockProvider).toHaveBeenCalledWith(
+      expect.objectContaining({ value: new Map() })
+    );
+  });
+});
diff --git a/plugins/ui/src/js/src/layout/ReactPanel.test.tsx b/plugins/ui/src/js/src/layout/ReactPanel.test.tsx
index 52f21857c..2e6b15283 100644
--- a/plugins/ui/src/js/src/layout/ReactPanel.test.tsx
+++ b/plugins/ui/src/js/src/layout/ReactPanel.test.tsx
@@ -7,9 +7,9 @@ import {
   ReactPanelManagerContext,
 } from './ReactPanelManager';
 import { ReactPanelProps } from './LayoutUtils';
+import PortalPanelManager from './PortalPanelManager';
 
 const mockPanelId = 'test-panel-id';
-jest.mock('shortid', () => jest.fn(() => mockPanelId));
 
 beforeEach(() => {
   jest.clearAllMocks();
@@ -20,18 +20,22 @@ function makeReactPanel({
   metadata = { name: 'test-name', type: 'test-type' },
   onClose = jest.fn(),
   onOpen = jest.fn(),
+  getPanelId = jest.fn(() => mockPanelId),
   title = 'test title',
 }: Partial<ReactPanelProps> & Partial<ReactPanelManager> = {}) {
   return (
-    <ReactPanelManagerContext.Provider
-      value={{
-        metadata,
-        onClose,
-        onOpen,
-      }}
-    >
-      <ReactPanel title={title}>{children}</ReactPanel>
-    </ReactPanelManagerContext.Provider>
+    <PortalPanelManager>
+      <ReactPanelManagerContext.Provider
+        value={{
+          getPanelId,
+          metadata,
+          onClose,
+          onOpen,
+        }}
+      >
+        <ReactPanel title={title}>{children}</ReactPanel>
+      </ReactPanelManagerContext.Provider>
+    </PortalPanelManager>
   );
 }
 
@@ -62,7 +66,7 @@ it('opens panel on mount, and closes panel on unmount', async () => {
 it('only calls open once if the panel has not closed and only children change', async () => {
   const onOpen = jest.fn();
   const onClose = jest.fn();
-  const metadata = { foo: 'bar' };
+  const metadata = { type: 'bar' };
   const children = 'hello';
   const { rerender } = render(
     makeReactPanel({ children, onOpen, onClose, metadata })
@@ -83,7 +87,7 @@ it('only calls open once if the panel has not closed and only children change',
 it('calls openComponent again after panel is closed only if the metadata changes', async () => {
   const onOpen = jest.fn();
   const onClose = jest.fn();
-  const metadata = { foo: 'bar' };
+  const metadata = { type: 'bar' };
   const children = 'hello';
   const { rerender } = render(
     makeReactPanel({
@@ -127,7 +131,7 @@ it('calls openComponent again after panel is closed only if the metadata changes
       children,
       onOpen,
       onClose,
-      metadata: { fiz: 'baz' },
+      metadata: { type: 'baz' },
     })
   );
 
diff --git a/plugins/ui/src/js/src/widget/WidgetHandler.test.tsx b/plugins/ui/src/js/src/widget/WidgetHandler.test.tsx
index cd12df8ad..d4827ded4 100644
--- a/plugins/ui/src/js/src/widget/WidgetHandler.test.tsx
+++ b/plugins/ui/src/js/src/widget/WidgetHandler.test.tsx
@@ -7,7 +7,6 @@ import {
   makeDocumentUpdatedJsonRpcString,
   makeWidget,
   makeWidgetDescriptor,
-  makeWidgetWrapper,
 } from './WidgetTestUtils';
 
 const mockApi = { Widget: { EVENT_MESSAGE: 'message' } };
@@ -24,10 +23,11 @@ jest.mock(
 );
 
 function makeWidgetHandler({
-  widget = makeWidgetWrapper(),
+  fetch = () => Promise.resolve(makeWidget()),
+  widget = makeWidgetDescriptor(),
   onClose = jest.fn(),
 }: Partial<WidgetHandlerProps> = {}) {
-  return <WidgetHandler widget={widget} onClose={onClose} />;
+  return <WidgetHandler fetch={fetch} widget={widget} onClose={onClose} />;
 }
 
 beforeEach(() => {
@@ -55,8 +55,8 @@ it('updates the document when event is received', async () => {
       makeDocumentUpdatedJsonRpcString(initialDocument)
     ),
   });
-  const wrapper = makeWidgetWrapper({ widget, fetch });
-  const { unmount } = render(makeWidgetHandler({ widget: wrapper }));
+
+  const { unmount } = render(makeWidgetHandler({ widget, fetch }));
   expect(fetch).toHaveBeenCalledTimes(1);
   expect(mockAddEventListener).not.toHaveBeenCalled();
   expect(mockDocumentHandler).not.toHaveBeenCalled();