From af6d6517a68587182071826329281ad521f09519 Mon Sep 17 00:00:00 2001
From: Jeremy Tuloup <jeremy.tuloup@gmail.com>
Date: Mon, 22 Aug 2022 11:04:46 +0200
Subject: [PATCH 1/2] Handle `KernelWidgetManager` in the JupyterLab
 `OutputModel`

---
 python/jupyterlab_widgets/src/output.ts | 24 ++++++++++++++++--------
 1 file changed, 16 insertions(+), 8 deletions(-)

diff --git a/python/jupyterlab_widgets/src/output.ts b/python/jupyterlab_widgets/src/output.ts
index 5f92fe2200..df2684faf9 100644
--- a/python/jupyterlab_widgets/src/output.ts
+++ b/python/jupyterlab_widgets/src/output.ts
@@ -7,7 +7,7 @@ import { JupyterLuminoPanelWidget } from '@jupyter-widgets/base';
 
 import { Panel } from '@lumino/widgets';
 
-import { WidgetManager } from './manager';
+import { KernelWidgetManager, WidgetManager } from './manager';
 
 import { OutputAreaModel, OutputArea } from '@jupyterlab/outputarea';
 
@@ -33,11 +33,14 @@ export class OutputModel extends outputBase.OutputModel {
       return false;
     };
 
-    this.widget_manager.context.sessionContext.kernelChanged.connect(
-      (sender, args) => {
-        this._handleKernelChanged(args);
-      }
-    );
+    // if the context is available, react on kernel changes
+    if (this.widget_manager instanceof WidgetManager) {
+      this.widget_manager.context.sessionContext.kernelChanged.connect(
+        (sender, args) => {
+          this._handleKernelChanged(args);
+        }
+      );
+    }
     this.listenTo(this, 'change:msg_id', this.reset_msg_id);
     this.listenTo(this, 'change:outputs', this.setOutputs);
     this.setOutputs();
@@ -60,7 +63,12 @@ export class OutputModel extends outputBase.OutputModel {
    * Reset the message id.
    */
   reset_msg_id(): void {
-    const kernel = this.widget_manager.context.sessionContext?.session?.kernel;
+    let kernel;
+    if (this.widget_manager instanceof WidgetManager) {
+      kernel = this.widget_manager.context.sessionContext?.session?.kernel;
+    } else if (this.widget_manager instanceof KernelWidgetManager) {
+      kernel = this.widget_manager.kernel;
+    }
     const msgId = this.get('msg_id');
     const oldMsgId = this.previous('msg_id');
 
@@ -114,7 +122,7 @@ export class OutputModel extends outputBase.OutputModel {
     }
   }
 
-  widget_manager: WidgetManager;
+  widget_manager: WidgetManager | KernelWidgetManager;
 
   private _msgHook: (msg: KernelMessage.IIOPubMessage) => boolean;
   private _outputs: OutputAreaModel;

From ed4ab3da8d9898b0bba4360e3804e5d8a097bd4d Mon Sep 17 00:00:00 2001
From: Jeremy Tuloup <jeremy.tuloup@gmail.com>
Date: Mon, 22 Aug 2022 11:19:20 +0200
Subject: [PATCH 2/2] Handle `LabWidgetManager`

---
 python/jupyterlab_widgets/src/output.ts | 11 +++--------
 1 file changed, 3 insertions(+), 8 deletions(-)

diff --git a/python/jupyterlab_widgets/src/output.ts b/python/jupyterlab_widgets/src/output.ts
index df2684faf9..37793bf193 100644
--- a/python/jupyterlab_widgets/src/output.ts
+++ b/python/jupyterlab_widgets/src/output.ts
@@ -7,7 +7,7 @@ import { JupyterLuminoPanelWidget } from '@jupyter-widgets/base';
 
 import { Panel } from '@lumino/widgets';
 
-import { KernelWidgetManager, WidgetManager } from './manager';
+import { LabWidgetManager, WidgetManager } from './manager';
 
 import { OutputAreaModel, OutputArea } from '@jupyterlab/outputarea';
 
@@ -63,12 +63,7 @@ export class OutputModel extends outputBase.OutputModel {
    * Reset the message id.
    */
   reset_msg_id(): void {
-    let kernel;
-    if (this.widget_manager instanceof WidgetManager) {
-      kernel = this.widget_manager.context.sessionContext?.session?.kernel;
-    } else if (this.widget_manager instanceof KernelWidgetManager) {
-      kernel = this.widget_manager.kernel;
-    }
+    const kernel = this.widget_manager.kernel;
     const msgId = this.get('msg_id');
     const oldMsgId = this.previous('msg_id');
 
@@ -122,7 +117,7 @@ export class OutputModel extends outputBase.OutputModel {
     }
   }
 
-  widget_manager: WidgetManager | KernelWidgetManager;
+  widget_manager: LabWidgetManager;
 
   private _msgHook: (msg: KernelMessage.IIOPubMessage) => boolean;
   private _outputs: OutputAreaModel;