fix: revive deephaven.ui widgets after a reconnect

plugins/ui/src/js/src/widget/WidgetHandler.test.tsx

@@ -5,6 +5,7 @@ import { type dh } from '@deephaven/jsapi-types';
 import { TestUtils } from '@deephaven/test-utils';
 import { type PluginModuleMap, PluginsContext } from '@deephaven/plugin';
 import { type Operation } from 'fast-json-patch';
+import { CALLABLE_KEY } from '../elements/utils/ElementUtils';
 import WidgetHandler, { type WidgetHandlerProps } from './WidgetHandler';
 import { type DocumentHandlerProps } from './DocumentHandler';
 import { type WidgetMessageEvent } from './WidgetTypes';
@@ -98,7 +99,8 @@ it('updates the document when event is received', async () => {
   const { unmount } = render(
     makeWidgetHandler({ widgetDescriptor: widget, initialData })
   );
-  expect(mockAddEventListener).toHaveBeenCalledTimes(1);
+  // One listener for 'message', one for 'reconnect'
+  expect(mockAddEventListener).toHaveBeenCalledTimes(2);
   expect(mockDocumentHandler).not.toHaveBeenCalled();
 
   // Verify setState was called with component state and appState
@@ -157,7 +159,97 @@ it('updates the document when event is received', async () => {
 
   expect(cleanup).not.toHaveBeenCalled();
   unmount();
-  expect(cleanup).toHaveBeenCalledTimes(1);
+  // Both the 'message' and 'reconnect' listeners are cleaned up
+  expect(cleanup).toHaveBeenCalledTimes(2);
+});
+
+it('resets the document and re-sends the last state when the widget reconnects', async () => {
+  const widget = makeWidgetDescriptor();
+  const mockAddEventListener = jest.fn((() =>
+    jest.fn()) as dh.Widget['addEventListener']);
+  const mockSendMessage = jest.fn();
+  const initialData = { state: { count: 0 } };
+  mockWidgetWrapper = {
+    widget: makeWidget({
+      addEventListener: mockAddEventListener,
+      getDataAsString: jest.fn(() => ''),
+      sendMessage: mockSendMessage,
+    }),
+    error: null,
+    api: jest.fn() as unknown as typeof dh,
+  };
+
+  const { unmount } = render(
+    makeWidgetHandler({ widgetDescriptor: widget, initialData })
+  );
+
+  const messageListener = mockAddEventListener.mock.calls.find(
+    call => call[0] === 'message'
+  )?.[1] as (event: WidgetMessageEvent) => void;
+  const reconnectListener = mockAddEventListener.mock.calls.find(
+    call => call[0] === 'reconnect'
+  )?.[1] as () => void;
+
+  // Respond to the initial setState, then send the initial document
+  // containing a callable, along with the latest widget state
+  await act(async () => {
+    messageListener(makeWidgetEventJsonRpcResponse(0));
+    messageListener(
+      makeWidgetEventDocumentPatched(
+        [
+          { op: 'add', path: '/foo', value: 'bar' },
+          { op: 'add', path: '/action', value: { [CALLABLE_KEY]: 'cb0' } },
+        ],
+        JSON.stringify({ count: 3 })
+      )
+    );
+  });
+
+  expect(mockDocumentHandler).toHaveBeenCalledTimes(1);
+  const staleDocument = mockDocumentHandler.mock.calls[0][0]
+    .children as unknown as {
+    foo: string;
+    action: unknown;
+  };
+  expect(typeof staleDocument.action).toBe('function');
+
+  mockSendMessage.mockClear();
+  mockDocumentHandler.mockClear();
+
+  // The same widget object fires 'reconnect' after reopening its stream
+  await act(async () => {
+    reconnectListener();
+  });
+
+  // setState is re-sent with the last state received from the server
+  expect(mockSendMessage).toHaveBeenCalledTimes(1);
+  const setStatePayload = JSON.parse(
+    mockSendMessage.mock.calls[0][0] as string
+  );
+  expect(setStatePayload.method).toBe('setState');
+  expect(setStatePayload.params[0]).toMatchObject({ count: 3 });
+
+  // The new stream renders from scratch: its from-empty patch must replace the
+  // stale document instead of merging with it, and callables must not be reused
+  await act(async () => {
+    messageListener(makeWidgetEventJsonRpcResponse(1));
+    messageListener(
+      makeWidgetEventDocumentPatched([
+        { op: 'add', path: '/action', value: { [CALLABLE_KEY]: 'cb0' } },
+      ])
+    );
+  });
+
+  expect(mockDocumentHandler).toHaveBeenCalledTimes(1);
+  const freshDocument = mockDocumentHandler.mock.calls[0][0]
+    .children as unknown as {
+    action: unknown;
+  };
+  expect(Object.keys(freshDocument)).toEqual(['action']);
+  expect(typeof freshDocument.action).toBe('function');
+  expect(freshDocument.action).not.toBe(staleDocument.action);
+
+  unmount();
 });
 
 it('updates the initial data only when widget has changed', async () => {
@@ -188,7 +280,7 @@ it('updates the initial data only when widget has changed', async () => {
       onClose,
     })
   );
-  expect(addEventListener).toHaveBeenCalledTimes(1);
+  expect(addEventListener).toHaveBeenCalledTimes(2);
   expect(mockDocumentHandler).not.toHaveBeenCalled();
   // Verify setState was called with component state and appState
   const setStatePayload1 = JSON.parse(sendMessage.mock.calls[0][0] as string);
@@ -256,8 +348,8 @@ it('updates the initial data only when widget has changed', async () => {
     })
   );
 
-  // Should have been called when the widget was updated
-  expect(cleanup).toHaveBeenCalledTimes(1);
+  // Should have been called for both listeners when the widget was updated
+  expect(cleanup).toHaveBeenCalledTimes(2);
   cleanup.mockClear();
 
   // eslint-disable-next-line prefer-destructuring
@@ -290,7 +382,8 @@ it('updates the initial data only when widget has changed', async () => {
 
   expect(cleanup).not.toHaveBeenCalled();
   unmount();
-  expect(cleanup).toHaveBeenCalledTimes(1);
+  // Both the 'message' and 'reconnect' listeners are cleaned up
+  expect(cleanup).toHaveBeenCalledTimes(2);
 });
 
 it('handles rendering widget error if widget is null (query disconnected)', async () => {
@@ -319,7 +412,7 @@ it('handles rendering widget error if widget is null (query disconnected)', asyn
       initialData: data1,
     })
   );
-  expect(mockAddEventListener).toHaveBeenCalledTimes(1);
+  expect(mockAddEventListener).toHaveBeenCalledTimes(2);
   expect(mockDocumentHandler).not.toHaveBeenCalled();
   const setStatePayloadErr = JSON.parse(sendMessage.mock.calls[0][0] as string);
   expect(setStatePayloadErr.method).toBe('setState');

plugins/ui/src/js/src/widget/WidgetHandler.tsx

@@ -148,6 +148,10 @@ function WidgetHandler({
     new Map<string, (...args: unknown[]) => void>()
   );
 
+  // The last state received from the server, so the widget can be restored to its
+  // current state (not just its initial one) after a reconnect
+  const lastReceivedStateRef = useRef<Record<string, unknown>>();
+
   // Bi-directional communication as defined in https://www.npmjs.com/package/json-rpc-2.0
   const jsonClient = useMemo(
     () =>
@@ -404,7 +408,7 @@ function WidgetHandler({
       log.debug('Adding methods to jsonClient');
       jsonClient.addMethod(
         METHOD_DOCUMENT_PATCHED,
-        async (params: [Operation[], string]) => {
+        async (params: [Operation[], string?]) => {
           log.debug2(METHOD_DOCUMENT_PATCHED, params);
           const [patch, stateParam] = params;
 
@@ -419,6 +423,16 @@ function WidgetHandler({
           if (stateParam != null) {
             try {
               const newState = JSON.parse(stateParam);
+              if (
+                newState == null ||
+                typeof newState !== 'object' ||
+                Array.isArray(newState)
+              ) {
+                throw new Error(
+                  `Expected state to be an object, received ${typeof newState}`
+                );
+              }
+              lastReceivedStateRef.current = newState;
               onDataChange({ state: newState });
             } catch (e) {
               log.warn(
@@ -526,6 +540,7 @@ function WidgetHandler({
       exportedObjectMap.current = widgetExportedObjectMap;
       exportedObjectCount.current = 0;
       renderedCallableMap.current.clear();
+      lastReceivedStateRef.current = undefined;
 
       // Set a var to the client that we know will not be null in the closure below
       const activeClient = jsonClient;
@@ -557,6 +572,32 @@ function WidgetHandler({
         }
       );
 
+      // After a disconnect+reconnect, the same widget object reopens its message stream and the
+      // server renders the widget from scratch: it sends a from-empty document patch with new
+      // callable and exported object ids. Drop all state tied to the old stream, then re-send
+      // the last known state so the server re-renders the document.
+      const reconnectCleanup = widget.addEventListener(
+        // This is defined as dh.Table.EVENT_RECONNECT in Core, use the constant value directly
+        // for the same reason as the 'message' listener above
+        'reconnect',
+        () => {
+          log.debug('Widget reconnected, resetting state', widget);
+          widgetExportedObjectMap.forEach(exportedObject => {
+            exportedObject.close();
+          });
+          widgetExportedObjectMap.clear();
+          exportedObjectCount.current = 0;
+          renderedCallableMap.current.clear();
+          // TODO: Remove unstable_batchedUpdates wrapper when upgrading to React 18
+          unstable_batchedUpdates(() => {
+            setInternalError(undefined);
+            setDocument(undefined);
+            setIsLoading(true);
+          });
+          sendSetState(lastReceivedStateRef.current ?? initialData?.state);
+        }
+      );
+
       log.debug('Receiving initial data');
       // We need to get the initial data and process it. If it's an old version of the plugin, it could be a documentUpdated command.
       receiveData(widget.getDataAsString(), widget.exportedObjects);
@@ -567,6 +608,7 @@ function WidgetHandler({
       return () => {
         log.debug('Cleaning up widget', widget);
         cleanup();
+        reconnectCleanup();
         widget.close();
 
         // Clean up any exported objects that haven't been closed yet

plugins/ui/src/js/src/widget/WidgetTestUtils.ts

@@ -10,15 +10,18 @@ import {
   type WidgetMessageEvent,
 } from './WidgetTypes';
 
-export function makeDocumentPatchedJsonRpc(patch: Operation[] = []): {
+export function makeDocumentPatchedJsonRpc(
+  patch: Operation[] = [],
+  state?: string
+): {
   jsonrpc: string;
   method: string;
-  params: [Operation[]];
+  params: [Operation[], string?];
 } {
   return {
     jsonrpc: '2.0',
     method: METHOD_DOCUMENT_PATCHED,
-    params: [patch],
+    params: state === undefined ? [patch] : [patch, state],
   };
 }
 
@@ -31,9 +34,10 @@ export function makeJsonRpcResponseString(id: number, result = ''): string {
 }
 
 export function makeDocumentPatchedJsonRpcString(
-  patch: Operation[] = []
+  patch: Operation[] = [],
+  state?: string
 ): string {
-  return JSON.stringify(makeDocumentPatchedJsonRpc(patch));
+  return JSON.stringify(makeDocumentPatchedJsonRpc(patch, state));
 }
 
 export function makeWidgetEvent(data = ''): WidgetMessageEvent {
@@ -54,9 +58,10 @@ export function makeWidgetEventJsonRpcResponse(
 }
 
 export function makeWidgetEventDocumentPatched(
-  patch: Operation[] = []
+  patch: Operation[] = [],
+  state?: string
 ): WidgetMessageEvent {
-  return makeWidgetEvent(makeDocumentPatchedJsonRpcString(patch));
+  return makeWidgetEvent(makeDocumentPatchedJsonRpcString(patch, state));
 }
 
 export function makeWidgetEventDocumentError(