From d7e9b9945aad3b3e02f3d8eee818864a08812e3e Mon Sep 17 00:00:00 2001
From: martinRenou <martin.renou@gmail.com>
Date: Tue, 9 Nov 2021 11:32:55 +0100
Subject: [PATCH] Add models-request message on the control comm channel

---
 ipywidgets/__init__.py                      |  2 +-
 ipywidgets/widgets/tests/test_send_state.py |  3 +-
 ipywidgets/widgets/widget.py                | 44 +++++++++++++--------
 packages/schema/messages.md                 |  4 +-
 4 files changed, 34 insertions(+), 19 deletions(-)

diff --git a/ipywidgets/__init__.py b/ipywidgets/__init__.py
index 7e905275bd..8ea8fd723d 100644
--- a/ipywidgets/__init__.py
+++ b/ipywidgets/__init__.py
@@ -38,7 +38,7 @@ def register_comm_target(kernel=None):
     if kernel is None:
         kernel = get_ipython().kernel
     kernel.comm_manager.register_target('jupyter.widget', Widget.handle_comm_opened)
-    kernel.comm_manager.register_target('jupyter.widget.control', Widget.handle_comm_opened_control)
+    kernel.comm_manager.register_target('jupyter.widget.control', Widget.handle_control_comm_opened)
 
 # deprecated alias
 handle_kernel = register_comm_target
diff --git a/ipywidgets/widgets/tests/test_send_state.py b/ipywidgets/widgets/tests/test_send_state.py
index 50cfb03076..07333b8ca3 100644
--- a/ipywidgets/widgets/tests/test_send_state.py
+++ b/ipywidgets/widgets/tests/test_send_state.py
@@ -29,5 +29,6 @@ def test_control():
     comm = DummyComm()
     Widget.close_all()
     w = SimpleWidget()
-    Widget.handle_comm_opened_control(comm, {})
+    Widget.handle_control_comm_opened(comm, {})
+    Widget.handle_control_comm_msg({'type': 'models-request'})
     assert comm.messages
diff --git a/ipywidgets/widgets/widget.py b/ipywidgets/widgets/widget.py
index d2c64c3889..492260931e 100644
--- a/ipywidgets/widgets/widget.py
+++ b/ipywidgets/widgets/widget.py
@@ -291,6 +291,7 @@ class Widget(LoggingHasTraits):
     # Class attributes
     #-------------------------------------------------------------------------
     _widget_construction_callback = None
+    _control_comm = None
 
     # widgets is a dictionary of all active widget objects
     widgets = {}
@@ -303,7 +304,6 @@ def close_all(cls):
         for widget in list(cls.widgets.values()):
             widget.close()
 
-
     @staticmethod
     def on_widget_constructed(callback):
         """Registers a callback to be called when a widget is constructed.
@@ -319,25 +319,37 @@ def _call_widget_constructed(widget):
             Widget._widget_construction_callback(widget)
 
     @classmethod
-    def handle_comm_opened_control(cls, comm, msg):
+    def handle_control_comm_opened(cls, comm, msg):
         version = msg.get('metadata', {}).get('version', '')
         if version.split('.')[0] != CONTROL_PROTOCOL_VERSION_MAJOR:
             raise ValueError("Incompatible widget control protocol versions: received version %r, expected version %r"%(version, __control_protocol_version__))
 
-        cls.get_manager_state()
-        widgets = Widget.widgets.values()
-        # build a single dict with the full widget state
-        full_state = {}
-        drop_defaults = False
-        for widget in widgets:
-            full_state[widget.model_id] = {
-                'model_name': widget._model_name,
-                'model_module': widget._model_module,
-                'model_module_version': widget._model_module_version,
-                'state': widget.get_state(drop_defaults=drop_defaults),
-            }
-        full_state, buffer_paths, buffers = _remove_buffers(full_state)
-        comm.send([full_state, buffer_paths], buffers=buffers)
+        cls._control_comm = comm
+        cls._control_comm.on_msg(cls.handle_control_comm_msg)
+
+    @classmethod
+    def handle_control_comm_msg(cls, msg):
+        # This shouldn't happen unless someone calls this method manually
+        if cls._control_comm is None:
+            raise RuntimeError('Control comm has not been properly opened')
+
+        if msg['content']['data']['type'] == 'models-request':
+            # Send back the full widgets state
+            cls.get_manager_state()
+            widgets = Widget.widgets.values()
+            full_state = {}
+            drop_defaults = False
+            for widget in widgets:
+                full_state[widget.model_id] = {
+                    'model_name': widget._model_name,
+                    'model_module': widget._model_module,
+                    'model_module_version': widget._model_module_version,
+                    'state': widget.get_state(drop_defaults=drop_defaults),
+                }
+            full_state, buffer_paths, buffers = _remove_buffers(full_state)
+            cls._control_comm.send([full_state, buffer_paths], buffers=buffers)
+        else:
+            raise RuntimeError('Control comm msg "{}" not implemented'.format(msg['content']['data']['type']))
 
     @staticmethod
     def handle_comm_opened(comm, msg):
diff --git a/packages/schema/messages.md b/packages/schema/messages.md
index 6f7085f7fe..8f0811798c 100644
--- a/packages/schema/messages.md
+++ b/packages/schema/messages.md
@@ -348,4 +348,6 @@ This is implemented in ipywidgets 7.7.
 
 ### The `jupyter.widget.control` comm target
 
-A kernel-side Jupyter widgets library registers a `jupyter.widget.control` comm target that is used for fetching all widgets states through a "one shot" comm message (one for all widget instances). The kernel-side widgets library must answer to the "comm-open" message with a comm message containing the full state of all widget instances.
+A kernel-side Jupyter widgets library registers a `jupyter.widget.control` comm target that is used for fetching all widgets states through a "one shot" comm message (one for all widget instances).
+
+The kernel-side widgets library must answer to the `{"type": "models-request"}` message with a comm message containing the full state of all widget instances.