diff --git a/.github/workflows/check-release.yml b/.github/workflows/check-release.yml
index 542206dc4..97b1f74d9 100644
--- a/.github/workflows/check-release.yml
+++ b/.github/workflows/check-release.yml
@@ -24,7 +24,7 @@ jobs:
         uses: jupyter-server/jupyter_releaser/.github/actions/check-release@v2
         with:
           token: ${{ secrets.GITHUB_TOKEN }}
-          version_spec: minor
+          version_spec: "12.34.56"
 
       - name: Upload Distributions
         uses: actions/upload-artifact@v4
diff --git a/packages/jupyter-ai/jupyter_ai/chat_handlers/base.py b/packages/jupyter-ai/jupyter_ai/chat_handlers/base.py
index 233099151..1ff5d13dd 100644
--- a/packages/jupyter-ai/jupyter_ai/chat_handlers/base.py
+++ b/packages/jupyter-ai/jupyter_ai/chat_handlers/base.py
@@ -182,7 +182,7 @@ def __init__(
         self.context_providers = context_providers
         self.message_interrupted = message_interrupted
         self.ychat = ychat
-        self.indexes_by_id: Dict[str, str] = {}
+        self.indexes_by_id: Dict[str, int] = {}
         """
         Indexes of messages in the YChat document by message ID.
 
@@ -282,17 +282,23 @@ async def _default_handle_exc(self, e: Exception, message: HumanChatMessage):
         )
         self.reply(response, message)
 
-    def write_message(self, body: str, id: Optional[str] = None) -> None:
-        """[Jupyter Chat only] Writes a message to the YChat shared document
-        that this chat handler is assigned to."""
+    def write_message(self, body: str, id: Optional[str] = None) -> str:
+        """
+        [Jupyter Chat only] Writes a message to the YChat shared document
+        that this chat handler is assigned to.
+
+        Returns the new message ID. This will be identical to the `id` argument
+        if passed.
+        """
+        # TODO: remove this once `ychat` becomes a required attribute.
         if not self.ychat:
-            return
+            return ""
 
         bot = self.ychat.get_user(BOT["username"])
         if not bot:
             self.ychat.set_user(BOT)
 
-        index = self.indexes_by_id.get(id, None)
+        index = self.indexes_by_id.get(id, None) if id else None
         id = id if id else str(uuid4())
         new_index = self.ychat.set_message(
             {
diff --git a/packages/jupyter-ai/jupyter_ai/extension.py b/packages/jupyter-ai/jupyter_ai/extension.py
index 7c02b9e40..83cc00c13 100644
--- a/packages/jupyter-ai/jupyter_ai/extension.py
+++ b/packages/jupyter-ai/jupyter_ai/extension.py
@@ -57,11 +57,11 @@
 JCOLLAB_VERSION = int(jupyter_collaboration_version[0])
 
 if JCOLLAB_VERSION >= 3:
-    from jupyter_server_ydoc.utils import (  # type:ignore[import-untyped]
+    from jupyter_server_ydoc.utils import (  # type:ignore[import-not-found,import-untyped]
         JUPYTER_COLLABORATION_EVENTS_URI,
     )
 else:
-    from jupyter_collaboration.utils import (  # type:ignore[import-untyped]
+    from jupyter_collaboration.utils import (  # type:ignore[import-not-found,import-untyped]
         JUPYTER_COLLABORATION_EVENTS_URI,
     )
 
@@ -294,6 +294,7 @@ async def get_chat(self, room_id: str) -> Optional[YChat]:
         if room_id in self.ychats_by_room:
             return self.ychats_by_room[room_id]
 
+        assert self.serverapp
         if JCOLLAB_VERSION >= 3:
             collaboration = self.serverapp.web_app.settings["jupyter_server_ydoc"]
             document = await collaboration.get_document(room_id=room_id, copy=False)
@@ -509,9 +510,11 @@ def _init_chat_handlers(
 
         TODO: Make `ychat` required once Jupyter Chat migration is complete.
         """
+        assert self.serverapp
+
         eps = entry_points()
         chat_handler_eps = eps.select(group="jupyter_ai.chat_handlers")
-        chat_handlers = {}
+        chat_handlers: Dict[str, BaseChatHandler] = {}
         chat_handler_kwargs = {
             "log": self.log,
             "config_manager": self.settings["jai_config_manager"],
diff --git a/packages/jupyter-ai/ui-tests/tests/jupyter-ai.spec.ts-snapshots/sidebar-linux.png b/packages/jupyter-ai/ui-tests/tests/jupyter-ai.spec.ts-snapshots/sidebar-linux.png
index 7b3103955..acfd60c1d 100644
Binary files a/packages/jupyter-ai/ui-tests/tests/jupyter-ai.spec.ts-snapshots/sidebar-linux.png and b/packages/jupyter-ai/ui-tests/tests/jupyter-ai.spec.ts-snapshots/sidebar-linux.png differ