Close #166. Cleanup orphaned widget models (i.e., models that are created as a side-effect of a output render) when they're no longer needed (i.e., the next time the output renders or if the session closes)

Author: cpsievert

js/src/comm.ts

@@ -96,7 +96,9 @@ export class ShinyComm {
     metadata?: any,
     buffers?: ArrayBuffer[] | ArrayBufferView[]
   ): string {
-    // I don't think we need to do anything here?
+    // The client-side model/comm has closed, let the server know so it can clean up
+    // the corresponding model/comm on its end.
+    Shiny.setInputValue("shinywidgets_comm_close", this.comm_id, {priority: "event"});
     return this.comm_id;
   }
 

js/src/output.ts

@@ -9,6 +9,7 @@ import type { ErrorsMessageValue } from 'rstudio-shiny/srcts/types/src/shiny/shi
  ******************************************************************************/
 
 class OutputManager extends HTMLManager {
+  orphaned_models: string[];
   // In a soon-to-be-released version of @jupyter-widgets/html-manager,
   // display_view()'s first "dummy" argument will be removed... this shim simply
   // makes it so that our manager can work with either version
@@ -91,20 +92,31 @@ class IPyWidgetOutput extends Shiny.OutputBinding {
     if (fill) el.classList.add("forward-fill-potential");
 
     // At this time point, we should've already handled an 'open' message, and so
-    // the model should be ready to use
+    // the model should already be registered
     const model = await manager.get_model(data.model_id);
     if (!model) {
       throw new Error(`No model found for id ${data.model_id}`);
     }
 
+    // Create a view and display it
     const view = await manager.create_view(model, {});
     await manager.display_view(view, {el: el});
 
-    // Don't allow more than one .lmWidget container, which can happen
-    // when the view is displayed more than once
-    // TODO: It's probably better to get view(s) from m.views and .remove() them
-    while (el.childNodes.length > 1) {
-      el.removeChild(el.childNodes[0]);
+    // If the model was orphaned, close it (and thus, the view as well) now
+    if (manager.orphaned_models.length > 0) {
+      for (const model_id of manager.orphaned_models) {
+        const model = await manager.get_model(model_id);
+        if (model) {
+          // Closing the model removes the previous view
+          await model.close();
+          // .close() isn't enough to remove manager's reference to it,
+          // and apparently the only way to remove it is through the `comm:close` event
+          // https://github.com/jupyter-widgets/ipywidgets/blob/303cae4/packages/base-manager/src/manager-base.ts#L330-L337
+          // https://github.com/jupyter-widgets/ipywidgets/blob/303cae4/packages/base/src/widget.ts#L251-L253
+          model.trigger("comm:close");
+        }
+      }
+      manager.orphaned_models = [];
     }
 
     // The ipywidgets container (.lmWidget)
@@ -189,21 +201,34 @@ Shiny.addCustomMessageHandler("shinywidgets_comm_open", (msg_txt) => {
 // Basically out version of https://github.com/jupyterlab/jupyterlab/blob/d33de15/packages/services/src/kernel/default.ts#L1200-L1215
 Shiny.addCustomMessageHandler("shinywidgets_comm_msg", (msg_txt) => {
   const msg = jsonParse(msg_txt);
-  manager.get_model(msg.content.comm_id).then(m => {
+  const id = msg.content.comm_id;
+  const model = manager.get_model(id);
+  if (!model) {
+    console.error(`Couldn't handle message for model ${id} because it doesn't exist.`);
+    return;
+  }
+  model.then(m => {
     // @ts-ignore for some reason IClassicComm doesn't have this method, but we do
     m.comm.handle_msg(msg);
   });
 });
 
-// TODO: test that this actually works
+
+// When `widget.close()` happens server-side, don't `.close()` client-side until the
+// next render. This is because we currently trigger a close _during_ the server-side
+// render, and thus, immediately closing the model removes the view before a new one is
+// ready.
+manager.orphaned_models = [];
 Shiny.addCustomMessageHandler("shinywidgets_comm_close", (msg_txt) => {
   const msg = jsonParse(msg_txt);
-  manager.get_model(msg.content.comm_id).then(m => {
-    // @ts-ignore for some reason IClassicComm doesn't have this method, but we do
-    m.comm.handle_close(msg)
-  });
+  manager.orphaned_models.push(msg.content.comm_id);
 });
 
+// At least currently, all widgets must be created within a session scope, so we can
+// clear the state (i.e., .close() all the widget models) when the session ends
+$(document).on("shiny:disconnected", () => {
+  manager.clear_state();
+});
 
 // Our version of https://github.com/jupyter-widgets/widget-cookiecutter/blob/9694718/%7B%7Bcookiecutter.github_project_name%7D%7D/js/lib/extension.js#L8
 function setBaseURL(x: string = '') {

shinywidgets/_render_widget_base.py

@@ -12,7 +12,6 @@
 from shiny import req
 from shiny.reactive._core import Context, get_current_context
 from shiny.render.renderer import Jsonifiable, Renderer, ValueFn
-from traitlets import Unicode
 
 from ._as_widget import as_widget
 from ._dependencies import widget_pkg
@@ -71,6 +70,9 @@ def __init__(
         self._contexts: set[Context] = set()
 
     async def render(self) -> Jsonifiable | None:
+        # Record the ids of any widgets created by executing the user function
+        widgets_before = set(WIDGET_INSTANCE_MAP.keys())
+
         value = await self.fn()
 
         # Attach value/widget attributes to user func so they can be accessed (in other reactive contexts)
@@ -89,16 +91,23 @@ async def render(self) -> Jsonifiable | None:
 
         self._widget = cast(WidgetT, widget)
 
+        # Get the widgets created within this render function
+        widgets_after = set(WIDGET_INSTANCE_MAP.keys())
+        widgets_diff = widgets_after - widgets_before
+
+        # When this render context gets invalidated, close widgets introduced by the
+        # user function. This _could_ be problematic if widget object get hoisted out
+        # of the function and used elsewhere, but if you're doing that, you're likely
+        # to have other problems anyway.
+        get_current_context().on_invalidate(lambda: close_widgets(widgets_diff))
+
         # Don't actually display anything unless this is a DOMWidget
         if not isinstance(widget, DOMWidget):
             return None
 
         return {
             "model_id": str(
-                cast(
-                    Unicode,
-                    widget.model_id,  # pyright: ignore[reportUnknownMemberType]
-                )
+                cast(str, widget.model_id)  # pyright: ignore[reportUnknownMemberType]
             ),
             "fill": fill,
         }
@@ -217,3 +226,17 @@ def set_layout_defaults(widget: Widget) -> Tuple[Widget, bool]:
             widget.chart = chart
 
     return (widget, fill)
+
+
+# Dictionary of all "active" widgets (ipywidgets automatically adds to this dictionary as
+# new widgets are created, but they won't get removed until the widget is explictly closed)
+WIDGET_INSTANCE_MAP = cast(dict[str, Widget], Widget.widgets)  # pyright: ignore[reportUnknownMemberType]
+
+# Given a set of widget ids, close those widgets. Closing should not only
+# remove the widget reference, but also send a message to the client to remove
+# the corresponding model from the widget manager.
+def close_widgets(ids: set[str]):
+    for id in ids:
+        widget = WIDGET_INSTANCE_MAP.get(id)
+        if widget is not None:
+            widget.close()

shinywidgets/_shinywidgets.py

@@ -26,6 +26,7 @@
 from ._cdn import SHINYWIDGETS_CDN_ONLY, SHINYWIDGETS_EXTENSION_WARNING
 from ._comm import BufferType, ShinyComm, ShinyCommManager
 from ._dependencies import require_dependency
+from ._render_widget_base import WIDGET_INSTANCE_MAP
 from ._utils import is_instance_of_class, package_dir
 
 __all__ = (
@@ -122,14 +123,28 @@ def init_shiny_widget(w: Widget):
 
     # Handle messages from the client. Note that widgets like qgrid send client->server messages
     # to figure out things like what filter to be shown in the table.
-    @reactive.Effect
+    @reactive.effect
     @reactive.event(session.input.shinywidgets_comm_send)
     def _():
         msg_txt = session.input.shinywidgets_comm_send()
         msg = json.loads(msg_txt)
         comm_id = msg["content"]["comm_id"]
-        comm: ShinyComm = COMM_MANAGER.comms[comm_id]
-        comm.handle_msg(msg)
+        if comm_id in COMM_MANAGER.comms:
+            comm: ShinyComm = COMM_MANAGER.comms[comm_id]
+            comm.handle_msg(msg)
+
+    # Handle a close message from the client.
+    @reactive.effect
+    @reactive.event(session.input.shinywidgets_comm_close)
+    def _():
+        comm_id = session.input.shinywidgets_comm_close()
+        # It seems the "proper" thing to do here is to .handle_close() on the comm
+        # but in practice, if we do that, ipywidgets will keep a reference to the
+        # widget instance via this widgets property, so it won't be garbage collected,
+        # unless we also call .close() on the widget instance.
+        w_obj = WIDGET_INSTANCE_MAP.get(comm_id)
+        if w_obj:
+            w_obj.close()
 
     def _restore_state():
         if old_comm_klass is not None: