feat: Return callables from callables in Deephaven UI (#540)

Fixes #322. Will be required for full completion of #321.

Currently a little hard to test, but this should demonstrate it if you
watch the JS debug logs on `WidgetHandler`.

Run this Python code.

- The `Normal` button should behave the same as latest dh.ui.
- The `Nested` button will return a callable as the direct result of the
callback function. This should allow us to do `const fn = await
callable(args);` basically. It is not called, but there is a log from
`WidgetHandler` after 5-10s that it was cleaned up. Or you can use the
`Memory` tab in dev tools and click the broom icon to trigger a manual
GC.
- The `Double Nested` button will return and wrap an object containing a
new callable. You can filter the browser logs for `WidgetUtils` and look
for `callable ID result string` which won't have the parsed function,
but will contain the object representing it as well as the other parts
of the returned object.
- The `Not serializable` button will basically do nothing except log an
error in the Python console. The returned result will include that it
was a serialization error to JS, but we do nothing with it.

```py
from deephaven import ui

@ui.component
def my_comp():
    on_press = ui.use_callback(lambda d: print(d), [])
    on_press_nested = ui.use_callback(lambda: print, [])
    on_press_double_nested = ui.use_callback(lambda: { "nestedFn": print, "some_val": 4 }, [])
    on_press_unserializable = ui.use_callback(lambda: set([1, 2, 3]), [])
    return [
        ui.button("Normal", on_press=on_press),
        ui.button("Nested", on_press=on_press_nested),
        ui.button("Double Nested", on_press=on_press_double_nested),
        ui.button("Not Serializable", on_press=on_press_unserializable)
    ]

c = my_comp()
```

plugins/ui/src/deephaven/ui/_internal/utils.py

@@ -127,7 +127,7 @@ def _wrapped_callable(
     func: Callable,
     *args: Any,
     **kwargs: Any,
-) -> None:
+) -> Any:
     """
     Filter the args and kwargs and call the specified function with the filtered args and kwargs.
 
@@ -139,14 +139,17 @@ def _wrapped_callable(
         func: The function to call
         *args: args, used by the dispatcher
         **kwargs: kwargs, used by the dispatcher
+
+    Returns:
+        The result of the function call.
     """
     args = args if max_args is None else args[:max_args]
     kwargs = (
         kwargs
         if kwargs_set is None
         else {k: v for k, v in kwargs.items() if k in kwargs_set}
     )
-    func(*args, **kwargs)
+    return func(*args, **kwargs)
 
 
 def wrap_callable(func: Callable) -> Callable:

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

@@ -19,6 +19,7 @@
 from .._internal import wrap_callable
 from ..elements import Element
 from ..renderer import NodeEncoder, Renderer, RenderedNode
+from ..renderer.NodeEncoder import CALLABLE_KEY
 from .._internal import RenderContext, StateUpdateCallable, ExportedRenderState
 from .ErrorCode import ErrorCode
 
@@ -98,6 +99,24 @@ class ElementMessageStream(MessageStream):
     Callables and render functions to be called on the next render loop.
     """
 
+    _callable_dict: dict[str, Callable]
+    """
+    Dict of callable IDs to callables.
+    This is intended to be used by the renderer and can be replaced on each render.
+    """
+
+    _temp_callable_dict: dict[str, Callable]
+    """
+    Dict of callable IDs to callables returned from other callables.
+    These are not generated by the renderer and can be removed by the client.
+    This should not be cleaned out on each render like _callable_dict.
+    """
+
+    _next_temp_callable_id: int
+    """
+    The next ID to use for temporary callables.
+    """
+
     _render_lock: threading.Lock
     """
     Lock to ensure only one thread is rendering at a time.
@@ -147,6 +166,9 @@ def __init__(self, element: Element, connection: MessageStream):
         self._renderer = Renderer(self._context)
         self._update_queue = Queue()
         self._callable_queue = Queue()
+        self._callable_dict = {}
+        self._temp_callable_dict = {}
+        self._next_temp_callable_id = 0
         self._render_lock = threading.Lock()
         self._is_dirty = False
         self._render_state = _RenderState.IDLE
@@ -327,6 +349,8 @@ def _make_request(self, method: str, *params: Any) -> dict[str, Any]:
     def _make_dispatcher(self) -> Dispatcher:
         dispatcher = Dispatcher()
         dispatcher["setState"] = self._set_state
+        dispatcher["callCallable"] = self._call_callable
+        dispatcher["closeCallable"] = self._close_callable
         return dispatcher
 
     def _set_state(self, state: ExportedRenderState) -> None:
@@ -340,6 +364,58 @@ def _set_state(self, state: ExportedRenderState) -> None:
         self._context.import_state(state)
         self._mark_dirty()
 
+    def _call_callable(self, callable_id: str, args: Any) -> Any:
+        """
+        Call a callable by its ID.
+        If the result is a callable, it is registered as a temporary callable.
+
+        Args:
+            callable_id: The ID of the callable to call
+            args: The array of arguments to pass to the callable. These will be spread as positional args to the callable.
+        """
+        logger.debug("Calling callable %s with %s", callable_id, args)
+        fn = self._callable_dict.get(callable_id) or self._temp_callable_dict.get(
+            callable_id
+        )
+        if fn is None:
+            logger.error("Callable not found: %s", callable_id)
+            return
+        result = fn(*args)
+
+        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:
+            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"
+            }
+
+    def _close_callable(self, callable_id: str) -> None:
+        """
+        Close a callable by its ID.
+
+        Args:
+            callable_id: The ID of the callable to close
+        """
+        logger.debug("Closing callable %s", callable_id)
+        self._callable_dict.pop(callable_id, None)
+        self._temp_callable_dict.pop(callable_id, None)
+
     def _send_document_update(
         self, root: RenderedNode, state: ExportedRenderState
     ) -> None:
@@ -366,11 +442,11 @@ def _send_document_update(
         payload = json.dumps(request)
         logger.debug(f"Sending payload: {payload}")
 
-        dispatcher = self._make_dispatcher()
+        callable_dict = {}
         for callable, callable_id in callable_id_dict.items():
             logger.debug("Registering callable %s", callable_id)
-            dispatcher[callable_id] = wrap_callable(callable)
-        self._dispatcher = dispatcher
+            callable_dict[callable_id] = wrap_callable(callable)
+        self._callable_dict = callable_dict
         if self._is_closed:
             # The connection is closed, so this component will not update anymore
             # delete the context so the objects in the collected scope are released

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

@@ -18,7 +18,7 @@ import { WidgetDescriptor } from '@deephaven/dashboard';
 import { useWidget } from '@deephaven/jsapi-bootstrap';
 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,
@@ -35,7 +35,7 @@ import {
   METHOD_DOCUMENT_UPDATED,
 } from './WidgetTypes';
 import DocumentHandler from './DocumentHandler';
-import { getComponentForElement } from './WidgetUtils';
+import { getComponentForElement, wrapCallable } from './WidgetUtils';
 import WidgetErrorView from './WidgetErrorView';
 import ReactPanelContentOverlayContext from '../layout/ReactPanelContentOverlayContext';
 
@@ -116,6 +116,15 @@ function WidgetHandler({
     [jsonClient]
   );
 
+  const callableFinalizationRegistry = useMemo(
+    () =>
+      new FinalizationRegistry(callableId => {
+        log.debug2('Closing callable', callableId);
+        jsonClient?.request('closeCallable', [callableId]);
+      }),
+    [jsonClient]
+  );
+
   const parseDocument = useCallback(
     /**
      * Parse the data from the server, replacing some of the nodes on the way.
@@ -127,6 +136,7 @@ function WidgetHandler({
      * @returns The parsed data
      */
     (data: string) => {
+      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);
@@ -136,10 +146,11 @@ function WidgetHandler({
         if (isCallableNode(value)) {
           const callableId = value[CALLABLE_KEY];
           log.debug2('Registering callableId', callableId);
-          return async (...args: unknown[]) => {
-            log.debug('Callable called', callableId, ...args);
-            return jsonClient?.request(callableId, args);
-          };
+          return wrapCallable(
+            jsonClient,
+            callableId,
+            callableFinalizationRegistry
+          );
         }
         if (isObjectNode(value)) {
           // Replace this node with the exported object
@@ -183,7 +194,7 @@ function WidgetHandler({
       );
       return parsedData;
     },
-    [jsonClient]
+    [jsonClient, callableFinalizationRegistry]
   );
 
   const updateExportedObjects = useCallback(

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

@@ -1,19 +1,40 @@
 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/model/ElementConstants';
-import { ElementNode, ELEMENT_KEY } from '../elements/utils/ElementUtils';
+import {
+  ElementNode,
+  ELEMENT_KEY,
+  CALLABLE_KEY,
+} from '../elements/utils/ElementUtils';
 import HTMLElementView from '../elements/HTMLElementView';
 import IconElementView from '../elements/IconElementView';
 import {
   elementComponentMap,
   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);
+  });
+});