diff --git a/code_setup.ipynb b/code_setup.ipynb
index 1d7947a..06c347e 100644
--- a/code_setup.ipynb
+++ b/code_setup.ipynb
@@ -17,8 +17,7 @@
    "source": [
     "# AiiDAlab Code Setup\n",
     "\n",
-    "This page provides an interface for creating AiiDA `Code` references to executable on \n",
-    "an AiiDA  `Computer` node."
+    "This page provides an interface for creating AiiDA codes on local or remote machines."
    ]
   },
   {
@@ -122,7 +121,15 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "widget.quick_setup_button.on_click(lambda _: render())"
+    "def on_code_setup_message_change(change):\n",
+    "    if \"created\" in change[\"new\"]:\n",
+    "        render()\n",
+    "\n",
+    "\n",
+    "widget.aiida_code_setup.observe(\n",
+    "    on_code_setup_message_change,\n",
+    "    \"message\",\n",
+    ")"
    ]
   }
  ],