diff --git a/code_setup.ipynb b/code_setup.ipynb
index 1d7947a..efd17e2 100644
--- a/code_setup.ipynb
+++ b/code_setup.ipynb
@@ -122,7 +122,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",
+    ")"
    ]
   }
  ],