Create new comm for existing models when restoring from kernel.

Author: Alan Fleming

packages/base-manager/src/manager-base.ts

@@ -499,6 +499,9 @@ export abstract class ManagerBase implements IWidgetManager {
               model.constructor as typeof WidgetModel
             )._deserialize_state(state.state, this);
             model!.set_state(deserializedState);
+            if (!model.comm_live) {
+              model.comm = await this._create_comm('jupyter.widget', widget_id);
+            }
           }
         } catch (error) {
           // Failed to create a widget model, we continue creating other models so that
@@ -750,12 +753,12 @@ export abstract class ManagerBase implements IWidgetManager {
 
   /**
    * Disconnect the widget manager from the kernel, setting each model's comm
-   * as dead.
+   * as undefined.
    */
   disconnect(): void {
     Object.keys(this._models).forEach((i) => {
       this._models[i].then((model) => {
-        model.comm_live = false;
+        model.comm = undefined;
       });
     });
   }

packages/base/src/errorwidget.ts

@@ -24,7 +24,9 @@ export function createErrorWidgetModel(
         error: error,
       };
       super(attributes, options);
-      this.comm_live = true;
+    }
+    get comm_live(): boolean {
+      return true;
     }
   }
   return ErrorWidget;

packages/base/src/widget.ts

@@ -160,7 +160,7 @@ export class WidgetModel extends Backbone.Model {
     // Attributes should be initialized here, since user initialization may depend on it
     this.widget_manager = options.widget_manager;
     this.model_id = options.model_id;
-    const comm = options.comm;
+    this.comm = options.comm;
 
     this.views = Object.create(null);
     this.state_change = Promise.resolve();
@@ -174,27 +174,23 @@ export class WidgetModel extends Backbone.Model {
     // _buffered_state_diff must be created *after* the super.initialize
     // call above. See the note in the set() method below.
     this._buffered_state_diff = {};
+  }
 
-    if (comm) {
-      // Remember comm associated with the model.
-      this.comm = comm;
+  get comm() {
+    return this._comm;
+  }
 
-      // Hook comm messages up to model.
+  set comm(comm: IClassicComm | undefined) {
+    this._comm = comm;
+    if (comm) {
       comm.on_close(this._handle_comm_closed.bind(this));
       comm.on_msg(this._handle_comm_msg.bind(this));
-
-      this.comm_live = true;
-    } else {
-      this.comm_live = false;
     }
+    this.trigger('comm_live_update');
   }
 
-  get comm_live(): boolean {
-    return this._comm_live;
-  }
-  set comm_live(x) {
-    this._comm_live = x;
-    this.trigger('comm_live_update');
+  get comm_live() {
+    return Boolean(this.comm);
   }
 
   /**
@@ -218,42 +214,42 @@ export class WidgetModel extends Backbone.Model {
    *
    * @returns - a promise that is fulfilled when all the associated views have been removed.
    */
-  close(comm_closed = false): Promise<void> {
+  async close(comm_closed = false): Promise<void> {
     // can only be closed once.
     if (this._closed) {
-      return Promise.resolve();
+      return;
     }
     this._closed = true;
-    if (this.comm && !comm_closed && this.comm_live) {
+    if (this._comm && !comm_closed) {
       try {
-        this.comm.close();
+        this._comm.close();
       } catch (err) {
         // Do Nothing
       }
     }
     this.stopListening();
     this.trigger('destroy', this);
-    if (this.comm) {
-      delete this.comm;
-    }
+    delete this._comm;
+
     // Delete all views of this model
     if (this.views) {
       const views = Object.keys(this.views).map((id: string) => {
         return this.views![id].then((view) => view.remove());
       });
       delete this.views;
-      return Promise.all(views).then(() => {
-        return;
-      });
+      await Promise.all(views);
+      return;
     }
-    return Promise.resolve();
   }
 
   /**
    * Handle when a widget comm is closed.
    */
   _handle_comm_closed(msg: KernelMessage.ICommCloseMsg): void {
-    this.comm_live = false;
+    if (!this.comm) {
+      return;
+    }
+    this.comm = undefined;
     this.trigger('comm:close');
     if (!this._closed) {
       this.close(true);
@@ -642,7 +638,7 @@ export class WidgetModel extends Backbone.Model {
    * This invokes a Backbone.Sync.
    */
   save_changes(callbacks?: {}): void {
-    if (this.comm_live) {
+    if (this.comm) {
       const options: any = { patch: true };
       if (callbacks) {
         options.callbacks = callbacks;
@@ -728,11 +724,9 @@ export class WidgetModel extends Backbone.Model {
   model_id: string;
   views?: { [key: string]: Promise<WidgetView> };
   state_change: Promise<any>;
-  comm?: IClassicComm;
   name: string;
   module: string;
-
-  private _comm_live: boolean;
+  private _comm?: IClassicComm;
   private _closed: boolean;
   private _state_lock: any;
   private _buffered_state_diff: any;

packages/base/test/src/widget_test.ts

@@ -194,15 +194,15 @@ describe('WidgetModel', function () {
       });
     });
 
-    it('sets the widget_manager, id, comm, and comm_live properties', function () {
-      const widgetDead = new WidgetModel({}, {
-        model_id: 'widgetDead',
+    it('sets the widget_manager, id, comm, properties', function () {
+      const widgetNoComm = new WidgetModel({}, {
+        model_id: 'noComm',
         widget_manager: this.manager,
       } as IBackboneModelOptions);
-      expect(widgetDead.model_id).to.equal('widgetDead');
-      expect(widgetDead.widget_manager).to.equal(this.manager);
-      expect(widgetDead.comm).to.be.undefined;
-      expect(widgetDead.comm_live).to.be.false;
+      expect(widgetNoComm.model_id).to.equal('noComm');
+      expect(widgetNoComm.widget_manager).to.equal(this.manager);
+      expect(widgetNoComm.comm).to.be.undefined;
+      expect(widgetNoComm.comm_live).to.be.false;
 
       const comm = new MockComm();
       const widgetLive = new WidgetModel({}, {

python/jupyterlab_widgets/src/manager.ts

@@ -302,7 +302,7 @@ abstract class LabWidgetManager extends ManagerBase implements IDisposable {
   get_state_sync(options: IStateOptions = {}): ReadonlyPartialJSONValue {
     const models = [];
     for (const model of this._modelsSync.values()) {
-      if (model.comm_live) {
+      if (model.comm) {
         models.push(model);
       }
     }
@@ -488,6 +488,7 @@ export class KernelWidgetManager extends LabWidgetManager {
       case 'restarting':
       case 'dead':
         this.disconnect();
+        this.clear_state();
         break;
     }
   }
@@ -510,7 +511,6 @@ export class KernelWidgetManager extends LabWidgetManager {
     this._restoredStatus = false;
     this._kernelRestoreInProgress = true;
     try {
-      await this.clear_state();
       await this._loadFromKernel();
     } catch {
       /* empty */