Fix appending rich widgets into an output (#16891)

DonJayamanne · web-flow · commit 9394edbe1c63 · 2025-08-08T07:40:14.000+02:00
* Fix appending rich widgets into an output

* Updates
diff --git a/src/test/datascience/notebook/helper.ts b/src/test/datascience/notebook/helper.ts
@@ -834,10 +834,10 @@ export async function waitForKernelToGetAutoSelectedImpl(
     return waitForKernelToChange(searchCriteria, notebookEditor!, timeout, skipAutoSelection);
 }
 
-const prewarmNotebooksDone = { done: false };
+let prewarmNotebooksDone: { ipyWidgetVersion: 7 | 8 | undefined } | undefined = undefined;
 export async function prewarmNotebooks() {
-    if (prewarmNotebooksDone.done) {
-        return;
+    if (prewarmNotebooksDone) {
+        return prewarmNotebooksDone;
     }
     const { serviceContainer } = await getServices();
     await closeActiveWindows();
@@ -855,11 +855,14 @@ export async function prewarmNotebooks() {
         const cell = window.activeNotebookEditor!.notebook.cellAt(0)!;
         logger.ci(`Running all cells in prewarm notebooks`);
         await Promise.all([waitForExecutionCompletedSuccessfully(cell, 60_000), runAllCellsInActiveNotebook()]);
+        const kernel = serviceContainer.get<IKernelProvider>(IKernelProvider).get(notebookEditor.notebook);
+        const ipyWidgetVersion = kernel?.ipywidgetsVersion;
         await closeActiveWindows();
         await shutdownAllNotebooks();
+        prewarmNotebooksDone = { ipyWidgetVersion };
+        return prewarmNotebooksDone;
     } finally {
         disposables.forEach((d) => d.dispose());
-        prewarmNotebooksDone.done = true;
     }
 }
 
diff --git a/src/test/datascience/widgets/notebooks/append_display_data_async.ipynb b/src/test/datascience/widgets/notebooks/append_display_data_async.ipynb
@@ -0,0 +1,30 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a0f3c2c2",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import asyncio\n",
+    "from IPython.display import display\n",
+    "from ipywidgets import HTML, Output\n",
+    "outputs = [Output() for _ in range(10)]\n",
+    "display(*outputs)\n",
+    "async def reproduce_bug():\n",
+    "    for i in range(10):\n",
+    "        await asyncio.sleep(0.01)\n",
+    "        outputs[i].append_display_data(HTML(f\"Content {i}\"))\n",
+    "asyncio.create_task(reproduce_bug())"
+   ]
+  }
+ ],
+ "metadata": {
+  "language_info": {
+   "name": "python"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}
diff --git a/src/test/datascience/widgets/notebooks/append_display_data_sync.ipynb b/src/test/datascience/widgets/notebooks/append_display_data_sync.ipynb
@@ -0,0 +1,34 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "0b12e473",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from IPython.display import display\n",
+    "from ipywidgets import HTML, Output\n",
+    "outputs = [Output() for _ in range(2)]\n",
+    "display(*outputs)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d40efbc0",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "outputs[0].append_display_data(HTML(\"Content 0\"))"
+   ]
+  }
+ ],
+ "metadata": {
+  "language_info": {
+   "name": "python"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}
diff --git a/src/test/datascience/widgets/standardWidgets.vscode.common.test.ts b/src/test/datascience/widgets/standardWidgets.vscode.common.test.ts
@@ -110,6 +110,7 @@ suite('Standard IPyWidget Tests @widgets', function () {
     this.retries(1);
     let editor: NotebookEditor;
     let comms: Utils;
+    let ipyWidgetVersion: 7 | 8 | undefined;
     suiteSetup(async function () {
         logger.info('Suite Setup Standard IPyWidget Tests');
         this.timeout(120_000);
@@ -120,7 +121,7 @@ suite('Standard IPyWidget Tests @widgets', function () {
         logger.info('Suite Setup Standard IPyWidget Tests, Step 3');
         await startJupyterServer();
         logger.info('Suite Setup Standard IPyWidget Tests, Step 4');
-        await prewarmNotebooks();
+        ({ ipyWidgetVersion } = await prewarmNotebooks());
         logger.info('Suite Setup Standard IPyWidget Tests, Step 5');
         sinon.restore();
         editor = (await createEmptyPythonNotebook(disposables, undefined, true)).editor;
@@ -593,4 +594,24 @@ suite('Standard IPyWidget Tests @widgets', function () {
             await assertOutputContainsHtml(cell, comms, ['Text Value is Bar']);
         });
     });
+    test('Sync append_display_data renders', async function () {
+        if (ipyWidgetVersion === 7) {
+            return this.skip();
+        }
+        await initializeNotebookForWidgetTest(disposables, { templateFile: 'append_display_data_sync.ipynb' }, editor);
+        const [cell1, cell2] = window.activeNotebookEditor!.notebook.getCells();
+        await executeCellAndWaitForOutput(cell1, comms);
+        await executeCellAndDontWaitForOutput(cell2);
+        await assertOutputContainsHtml(cell1, comms, ['Content 0'], '.widget-output');
+    });
+    test('Async append_display_data renders', async function () {
+        if (ipyWidgetVersion === 7) {
+            return this.skip();
+        }
+        await initializeNotebookForWidgetTest(disposables, { templateFile: 'append_display_data_async.ipynb' }, editor);
+        const [cell] = window.activeNotebookEditor!.notebook.getCells();
+        await executeCellAndWaitForOutput(cell, comms);
+        // Check that at least one output widget rendered content
+        await assertOutputContainsHtml(cell, comms, ['Content 3'], '.widget-output');
+    });
 });
diff --git a/src/webviews/extension-side/ipywidgets/rendererComms.ts b/src/webviews/extension-side/ipywidgets/rendererComms.ts
@@ -9,7 +9,7 @@ import { Disposable, NotebookDocument, NotebookEditor, NotebookRendererMessaging
 import { IKernel, IKernelProvider } from '../../../kernels/types';
 import { IControllerRegistration } from '../../../notebooks/controllers/types';
 import { IExtensionSyncActivationService } from '../../../platform/activation/types';
-import { WIDGET_MIMETYPE } from '../../../platform/common/constants';
+import { WIDGET_MIMETYPE, Identifiers } from '../../../platform/common/constants';
 import { dispose } from '../../../platform/common/utils/lifecycle';
 import { IDisposable } from '../../../platform/common/types';
 import { noop } from '../../../platform/common/utils/misc';
@@ -84,11 +84,23 @@ export class IPyWidgetRendererComms implements IExtensionSyncActivationService {
                 jupyterLab.KernelMessage.isExecuteResultMsg(msg)
             ) {
                 this.trackModelId(kernel.notebook, msg);
+            } else if (
+                jupyterLab.KernelMessage.isCommOpenMsg(msg) &&
+                // Track widget model ids as soon as the comm opens to avoid races with renderer queries.
+                msg.content?.target_name === Identifiers.DefaultCommTarget &&
+                typeof msg.content?.comm_id === 'string'
+            ) {
+                this.addModelId(kernel.notebook, msg.content.comm_id);
             }
         };
         iopubMessage.connect(handler);
         this.disposables.push(new Disposable(() => iopubMessage.disconnect(handler)));
     }
+    private addModelId(notebook: NotebookDocument, modelId: string) {
+        const set = this.widgetOutputsPerNotebook.get(notebook) || new Set<string>();
+        set.add(modelId);
+        this.widgetOutputsPerNotebook.set(notebook, set);
+    }
     private trackModelId(
         notebook: NotebookDocument,
         msg: {
@@ -101,9 +113,7 @@ export class IPyWidgetRendererComms implements IExtensionSyncActivationService {
         if (output.data && typeof output.data === 'object' && WIDGET_MIMETYPE in output.data) {
             const widgetData = output.data[WIDGET_MIMETYPE] as WidgetData;
             if (widgetData && 'model_id' in widgetData) {
-                const set = this.widgetOutputsPerNotebook.get(notebook) || new Set<string>();
-                set.add(widgetData.model_id);
-                this.widgetOutputsPerNotebook.set(notebook, set);
+                this.addModelId(notebook, widgetData.model_id);
             }
         }
     }
diff --git a/src/webviews/webview-side/ipywidgets/kernel/manager.ts b/src/webviews/webview-side/ipywidgets/kernel/manager.ts
@@ -18,7 +18,7 @@ import {
     ScriptLoader
 } from './types';
 import { KernelSocketOptions } from '../../../../kernels/types';
-import { Deferred, createDeferred } from '../../../../platform/common/utils/async';
+import { Deferred, createDeferred, raceTimeout } from '../../../../platform/common/utils/async';
 import { IInteractiveWindowMapping, IPyWidgetMessages, InteractiveWindowMessages } from '../../../../messageTypes';
 import { WIDGET_MIMETYPE, WIDGET_STATE_MIMETYPE } from '../../../../platform/common/constants';
 import { NotebookMetadata } from '../../../../platform/common/utils';
@@ -181,7 +181,13 @@ export class WidgetManager implements IIPyWidgetManager, IMessageHandler {
         // Check if we have processed the data for this model.
         // If not wait.
         if (!this.modelIdsToBeDisplayed.has(modelId)) {
-            this.modelIdsToBeDisplayed.set(modelId, createDeferred());
+            const deferred = createDeferred<void>();
+            this.modelIdsToBeDisplayed.set(modelId, deferred);
+            // Try to get the model if available.
+            const modelPromise = this.manager.get_model(modelId);
+            if (modelPromise && (await raceTimeout(1_000, modelPromise.catch(noop)))) {
+                deferred.resolve();
+            }
         }
         // Wait until it is flagged as ready to be processed.
         // This widget manager must have received this message and performed all operations before this.
@@ -315,17 +321,21 @@ export class WidgetManager implements IIPyWidgetManager, IMessageHandler {
 
         if (
             !jupyterLab.KernelMessage.isDisplayDataMsg(payload) &&
+            !jupyterLab.KernelMessage.isUpdateDisplayDataMsg?.(payload) &&
             !jupyterLab.KernelMessage.isExecuteResultMsg(payload)
         ) {
             return;
         }
-        const displayMsg = payload as KernelMessage.IDisplayDataMsg | KernelMessage.IExecuteResultMsg;
+        const displayMsg = payload as
+            | KernelMessage.IDisplayDataMsg
+            | KernelMessage.IUpdateDisplayDataMsg
+            | KernelMessage.IExecuteResultMsg;
 
         if (displayMsg.content && displayMsg.content.data && displayMsg.content.data[WIDGET_MIMETYPE]) {
             // eslint-disable-next-line @typescript-eslint/no-explicit-any
             const data = displayMsg.content.data[WIDGET_MIMETYPE] as any;
             const modelId = data.model_id;
-            logMessage(`Received display data message ${modelId}`);
+            logMessage(`Received display/update data message ${modelId}`);
             let deferred = this.modelIdsToBeDisplayed.get(modelId);
             if (!deferred) {
                 deferred = createDeferred();
@@ -337,8 +347,7 @@ export class WidgetManager implements IIPyWidgetManager, IMessageHandler {
             const modelPromise = this.manager.get_model(data.model_id);
             if (modelPromise) {
                 try {
-                    const m = await modelPromise;
-                    console.log(m);
+                    await modelPromise;
                     deferred.resolve();
                 } catch (ex) {
                     deferred.reject(ex);
