diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml
index 018fb8f..1ab3e45 100644
--- a/.github/workflows/build.yml
+++ b/.github/workflows/build.yml
@@ -151,6 +151,25 @@ jobs:
           ui-tests/test-results
           ui-tests/playwright-report
 
+  typing-tests:
+    name: Typing test
+    runs-on: ubuntu-latest
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v4
+
+      - name: Base Setup
+        uses: jupyterlab/maintainer-tools/.github/actions/base-setup@v1
+
+      - name: Install extension dependencies and build the extension
+        run: ./scripts/dev-install.sh
+
+      - name: Run mypy
+        run: |
+          set -eux
+          mypy --version
+          mypy python/jupyterlab-chat
+
   check_links:
     name: Check Links
     runs-on: ubuntu-latest
diff --git a/pyproject.toml b/pyproject.toml
index a465cbf..29e0e9f 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -54,3 +54,6 @@ before-bump-version = [
 
 [tool.check-wheel-contents]
 ignore = ["W002"]
+
+[tool.mypy]
+check_untyped_defs = true
diff --git a/python/jupyterlab-chat/jupyterlab_chat/__init__.py b/python/jupyterlab-chat/jupyterlab_chat/__init__.py
index e4855ea..08daf35 100644
--- a/python/jupyterlab-chat/jupyterlab_chat/__init__.py
+++ b/python/jupyterlab-chat/jupyterlab_chat/__init__.py
@@ -11,6 +11,8 @@
     warnings.warn("Importing 'jupyterlab_chat' outside a proper installation.")
     __version__ = "dev"
 
+from .ychat import YChat
+
 
 def _jupyter_labextension_paths():
     return [{
diff --git a/python/jupyterlab-chat/jupyterlab_chat/py.typed b/python/jupyterlab-chat/jupyterlab_chat/py.typed
new file mode 100644
index 0000000..e69de29
diff --git a/python/jupyterlab-chat/jupyterlab_chat/ychat.py b/python/jupyterlab-chat/jupyterlab_chat/ychat.py
index b3cdc17..519b53f 100644
--- a/python/jupyterlab-chat/jupyterlab_chat/ychat.py
+++ b/python/jupyterlab-chat/jupyterlab_chat/ychat.py
@@ -72,7 +72,7 @@ def get_users(self) -> dict[str, dict[str, str]]:
         Returns the users of the document.
         :return: Document's users.
         """
-        return self._yusers.to_py()
+        return self._yusers.to_py() or {}
 
     def set_user(self, user: dict[str, str]) -> None:
         """
@@ -95,7 +95,7 @@ def get_messages(self) -> list[dict]:
         Returns the messages of the document.
         :return: Document's messages.
         """
-        return self._ymessages.to_py()
+        return self._ymessages.to_py() or []
 
     def add_message(self, message: dict) -> int:
         """
@@ -119,19 +119,20 @@ def update_message(self, message: dict, index: int, append: bool = False):
                 message["body"] = initial_message["body"] + message["body"]
             self._ymessages.insert(index, message)
 
-    def set_message(self, message: dict, index: int | None = None, append: bool = False):
+    def set_message(self, message: dict, index: int | None = None, append: bool = False) -> int:
         """
         Update or append a message.
         """
 
+        initial_message: dict | None = None
         if index is not None and 0 <= index < len(self._ymessages):
-            initial_message = self._ymessages[index]
+            initial_message = self.get_messages()[index]
         else:
             return self.add_message(message)
 
         if not initial_message["id"] == message["id"]:
             initial_message, index = self.get_message(message["id"])
-            if initial_message is None:
+            if index is None:
                 return self.add_message(message)
 
         self.update_message(message, index, append)
@@ -147,7 +148,7 @@ def get_metadata(self) -> dict[str, dict]:
         """
         Returns the metadata of the document.
         """
-        return self._ymetadata.to_py()
+        return self._ymetadata.to_py() or {}
 
     def set_metadata(self, name: str, metadata: dict):
         """
@@ -164,7 +165,7 @@ async def create_id(self) -> str:
         self.set_id(id)
         return id
 
-    def get_id(self) -> str:
+    def get_id(self) -> str | None:
         """
         Returns the ID of the document.
         """
@@ -255,7 +256,7 @@ def _timestamp_new_messages(self, event: ArrayEvent) -> None:
         index = 0
         inserted_count = -1
         deleted_count = -1
-        for value in event.delta:
+        for value in event.delta:  # type:ignore[attr-defined]
             if "retain" in value.keys():
                 index = value["retain"]
             elif "insert" in value.keys():
@@ -268,7 +269,7 @@ def _timestamp_new_messages(self, event: ArrayEvent) -> None:
             return
 
         for idx in range(index, index + inserted_count):
-            message = self._ymessages[idx]
+            message = self.get_messages()[idx]
             if message and message.get("raw_time", True):
                 self.create_task(self._set_timestamp(idx, timestamp))
 
@@ -279,7 +280,7 @@ async def _set_timestamp(self, msg_idx: int, timestamp: float):
         with self._ydoc.transaction():
             # Remove the message from the list and modify the timestamp
             try:
-                message = self._ymessages[msg_idx]
+                message = self.get_messages()[msg_idx]
             except IndexError:
                 return
 
diff --git a/python/jupyterlab-chat/pyproject.toml b/python/jupyterlab-chat/pyproject.toml
index 3d28553..58146f9 100644
--- a/python/jupyterlab-chat/pyproject.toml
+++ b/python/jupyterlab-chat/pyproject.toml
@@ -37,6 +37,7 @@ dynamic = ["version", "description", "authors", "urls", "keywords"]
 [project.optional-dependencies]
 test = [
     "coverage",
+    "mypy",
     "pytest",
     "pytest-asyncio",
     "pytest-cov",