feat: Return callables from callables in Deephaven UI

plugins/ui/src/deephaven/ui/object_types/ElementMessageStream.py

@@ -113,7 +113,7 @@ class ElementMessageStream(MessageStream):
     This should not be cleaned out on each render like _callable_dict.
     """
 
-    _temp_callable_next_id: int
+    _next_temp_callable_id: int
     """
     The next ID to use for temporary callables.
     """
@@ -169,7 +169,7 @@ def __init__(self, element: Element, connection: MessageStream):
         self._callable_queue = Queue()
         self._callable_dict = {}
         self._temp_callable_dict = {}
-        self._temp_callable_next_id = 0
+        self._next_temp_callable_id = 0
         self._render_lock = threading.Lock()
         self._is_dirty = False
         self._render_state = _RenderState.IDLE
@@ -382,17 +382,26 @@ def _call_callable(self, callable_id: str, args: Any) -> Any:
             logger.error("Callable not found: %s", callable_id)
             return
         result = fn(*args)
-        if callable(result):
-            new_id = f"tempCb{self._temp_callable_next_id}"
-            self._temp_callable_next_id += 1
-            self._temp_callable_dict[new_id] = result
-            return {
-                CALLABLE_KEY: new_id,
-            }
+
+        def serialize_callables(node: Any) -> Any:
+            if callable(node):
+                new_id = f"tempCb{self._next_temp_callable_id}"
+                self._next_temp_callable_id += 1
+                self._temp_callable_dict[new_id] = node
+                return {
+                    CALLABLE_KEY: new_id,
+                }
+            raise TypeError(
+                f"A Deephaven UI callback returned a non-serializable value. Object of type {type(node).__name__} is not JSON serializable"
+            )
+
         try:
-            json.dumps(result)
-            return result
-        except Exception:
+            return json.dumps(result, default=serialize_callables)
+        except Exception as e:
+            # This is shown to the user in the Python console
+            # The stack trace from logger.exception is useless to the user
+            # Stack trace only includes the internals of the serialization process
+            logger.error(e)
             return {
                 "serialization_error": f"Cannot serialize callable {callable_id} result"
             }

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

@@ -17,7 +17,7 @@ import {
 import { WidgetDescriptor } from '@deephaven/dashboard';
 import type { dh } from '@deephaven/jsapi-types';
 import Log from '@deephaven/log';
-import { EMPTY_FUNCTION } from '@deephaven/utils';
+import { EMPTY_FUNCTION, assertNotNull } from '@deephaven/utils';
 import {
   CALLABLE_KEY,
   OBJECT_KEY,
@@ -133,10 +133,7 @@ function WidgetHandler({
      * @returns The parsed data
      */
     (data: string) => {
-      if (jsonClient == null) {
-        log.warn('No jsonClient set. Skipping parsing data');
-        return;
-      }
+      assertNotNull(jsonClient);
       // Keep track of exported objects that are no longer in use after this render.
       // We close those objects that are no longer referenced, as they will never be referenced again.
       const deadObjectMap = new Map(exportedObjectMap.current);

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

@@ -1,7 +1,13 @@
 import React from 'react';
+import type { JSONRPCServerAndClient } from 'json-rpc-2.0';
 import { Text } from '@deephaven/components';
+import { TestUtils } from '@deephaven/utils';
 import { ELEMENT_NAME, ELEMENT_PREFIX } from '../elements/ElementConstants';
-import { ElementNode, ELEMENT_KEY } from '../elements/ElementUtils';
+import {
+  ElementNode,
+  ELEMENT_KEY,
+  CALLABLE_KEY,
+} from '../elements/ElementUtils';
 import HTMLElementView from '../elements/HTMLElementView';
 import IconElementView from '../elements/IconElementView';
 import { SPECTRUM_ELEMENT_TYPE_PREFIX } from '../elements/SpectrumElementUtils';
@@ -11,8 +17,23 @@ import {
   getComponentForElement,
   getComponentTypeForElement,
   getPreservedData,
+  wrapCallable,
 } from './WidgetUtils';
 
+const mockJsonRequest = jest.fn(() =>
+  Promise.resolve(JSON.stringify({ result: 'mock' }))
+);
+
+const mockJsonClient = TestUtils.createMockProxy<JSONRPCServerAndClient>({
+  request: mockJsonRequest,
+});
+
+const mockFinilizationRegistry = TestUtils.createMockProxy<
+  FinalizationRegistry<unknown>
+>({
+  register: jest.fn(),
+});
+
 describe('getComponentTypeForElement', () => {
   it.each(
     Object.keys(elementComponentMap) as (keyof typeof elementComponentMap)[]
@@ -91,3 +112,102 @@ describe('getPreservedData', () => {
     expect(actual).toEqual({ panelIds: widgetData.panelIds });
   });
 });
+
+describe('wrapCallable', () => {
+  beforeEach(() => {
+    jest.clearAllMocks();
+  });
+
+  it('should return a function that sends a request to the client', () => {
+    wrapCallable(mockJsonClient, 'testMethod', mockFinilizationRegistry)();
+    expect(mockJsonClient.request).toHaveBeenCalledWith('callCallable', [
+      'testMethod',
+      [],
+    ]);
+  });
+
+  it('should return a function that sends a request to the client with args', () => {
+    wrapCallable(
+      mockJsonClient,
+      'testMethod',
+      mockFinilizationRegistry
+    )('a', { b: 'b' });
+    expect(mockJsonClient.request).toHaveBeenCalledWith('callCallable', [
+      'testMethod',
+      ['a', { b: 'b' }],
+    ]);
+  });
+
+  it('should register the function in the finalization registry', () => {
+    const wrapped = wrapCallable(
+      mockJsonClient,
+      'testMethod',
+      mockFinilizationRegistry
+    );
+
+    expect(mockFinilizationRegistry.register).toHaveBeenCalledWith(
+      wrapped,
+      'testMethod',
+      wrapped
+    );
+  });
+
+  it('should wrap returned callables', async () => {
+    mockJsonRequest.mockResolvedValueOnce(
+      JSON.stringify({
+        [CALLABLE_KEY]: 'nestedCb',
+      })
+    );
+
+    const wrappedResult = await wrapCallable(
+      mockJsonClient,
+      'testMethod',
+      mockFinilizationRegistry
+    )();
+    expect(wrappedResult).toBeInstanceOf(Function);
+
+    expect(mockFinilizationRegistry.register).toHaveBeenCalledTimes(2);
+    expect(mockFinilizationRegistry.register).toHaveBeenLastCalledWith(
+      wrappedResult,
+      'nestedCb',
+      wrappedResult
+    );
+  });
+
+  it('should wrap nested returned callables', async () => {
+    mockJsonRequest.mockResolvedValueOnce(
+      JSON.stringify({
+        nestedCallable: {
+          [CALLABLE_KEY]: 'nestedCb',
+        },
+        someOtherProp: 'mock',
+      })
+    );
+
+    const wrappedResult = (await wrapCallable(
+      mockJsonClient,
+      'testMethod',
+      mockFinilizationRegistry
+    )()) as { nestedCallable: () => void; someOtherProp: string };
+
+    expect(wrappedResult).toMatchObject({
+      nestedCallable: expect.any(Function),
+      someOtherProp: 'mock',
+    });
+
+    expect(mockFinilizationRegistry.register).toHaveBeenCalledTimes(2);
+    expect(mockFinilizationRegistry.register).toHaveBeenLastCalledWith(
+      wrappedResult.nestedCallable,
+      'nestedCb',
+      wrappedResult.nestedCallable
+    );
+  });
+
+  it('should reject if the result is not parseable', () => {
+    mockJsonRequest.mockResolvedValueOnce('not a json string');
+
+    expect(
+      wrapCallable(mockJsonClient, 'testMethod', mockFinilizationRegistry)()
+    ).rejects.toBeInstanceOf(Error);
+  });
+});

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

@@ -138,15 +138,32 @@ export function wrapCallable(
 ): (...args: unknown[]) => Promise<unknown> {
   const callable = async (...args: unknown[]) => {
     log.debug2(`Callable ${callableId} called`, args);
-    const result = await jsonClient.request('callCallable', [callableId, args]);
-    log.debug2(`Callable ${callableId} result`, result);
-    if (isCallableNode(result)) {
-      const nestedCallable = wrapCallable(
-        jsonClient,
-        result[CALLABLE_KEY],
-        registry
-      );
-      return nestedCallable;
+    const resultString = await jsonClient.request('callCallable', [
+      callableId,
+      args,
+    ]);
+
+    log.debug2(`Callable ${callableId} result string`, resultString);
+
+    try {
+      // Do NOT add anything that logs result
+      // It creates a strong ref to the result object in the console logs
+      // As a result, any returned callables will never be GC'd and the finalization registry will never clean them up
+      const result = JSON.parse(resultString, (key, value) => {
+        if (isCallableNode(value)) {
+          const nestedCallable = wrapCallable(
+            jsonClient,
+            value[CALLABLE_KEY],
+            registry
+          );
+          return nestedCallable;
+        }
+        return value;
+      });
+
+      return result;
+    } catch {
+      throw new Error(`Error parsing callable result: ${resultString}`);
     }
   };