diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml
index 5ba45bb4..656cf271 100644
--- a/.github/workflows/build.yml
+++ b/.github/workflows/build.yml
@@ -120,7 +120,7 @@ jobs:
         python -m pip install "jupyterlab>=4.0.0,<5" jupyterlab_chat*.whl
 
     - name: Install dependencies
-      working-directory: python/jupyterlab-chat/ui-tests
+      working-directory: ui-tests
       env:
         YARN_ENABLE_IMMUTABLE_INSTALLS: 0
         PLAYWRIGHT_SKIP_BROWSER_DOWNLOAD: 1
@@ -135,10 +135,10 @@ jobs:
 
     - name: Install browser
       run: jlpm playwright install chromium
-      working-directory: python/jupyterlab-chat/ui-tests
+      working-directory: ui-tests
 
     - name: Execute integration tests
-      working-directory: python/jupyterlab-chat/ui-tests
+      working-directory: ui-tests
       run: |
         jlpm playwright test --retries=2
 
@@ -148,8 +148,27 @@ jobs:
       with:
         name: jupyterlab_chat-playwright-tests
         path: |
-          python/jupyterlab-chat/ui-tests/test-results
-          python/jupyterlab-chat/ui-tests/playwright-report
+          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
diff --git a/.github/workflows/update-integration-tests.yml b/.github/workflows/update-integration-tests.yml
index 2ac7f1fe..aee19087 100644
--- a/.github/workflows/update-integration-tests.yml
+++ b/.github/workflows/update-integration-tests.yml
@@ -79,4 +79,4 @@ jobs:
           github_token: ${{ secrets.GITHUB_TOKEN }}
           # Playwright knows how to start JupyterLab server
           start_server_script: 'null'
-          test_folder: python/jupyterlab-chat/ui-tests
+          test_folder: ui-tests
diff --git a/.gitignore b/.gitignore
index 9405cc57..014691c3 100644
--- a/.gitignore
+++ b/.gitignore
@@ -18,8 +18,8 @@ packages/jupyter-chat/junit.xml
 packages/jupyterlab-chat/junit.xml
 
 # Integration tests
-python/jupyterlab-chat/ui-tests/test-results/
-python/jupyterlab-chat/ui-tests/playwright-report/
+ui-tests/test-results/
+ui-tests/playwright-report/
 
 # Created by https://www.gitignore.io/api/python
 # Edit at https://www.gitignore.io/?templates=python
diff --git a/.licenserc.yaml b/.licenserc.yaml
index a4b17bab..08084806 100644
--- a/.licenserc.yaml
+++ b/.licenserc.yaml
@@ -8,6 +8,7 @@ header:
       Distributed under the terms of the Modified BSD License.
 
   paths-ignore:
+    - '**/*.typed'
     - '**/*.ipynb'
     - '**/*.json'
     - '**/*.md'
diff --git a/docs/source/developers/contributing/jupyterlab-chat-extension.md b/docs/source/developers/contributing/jupyterlab-chat-extension.md
index 7b69ace4..71a4bccc 100644
--- a/docs/source/developers/contributing/jupyterlab-chat-extension.md
+++ b/docs/source/developers/contributing/jupyterlab-chat-extension.md
@@ -57,14 +57,14 @@ jlpm test
 
 ### Integration tests
 
-The integration tests are located in *python/jupyterlab-chat/ui-tests*.
+The integration tests are located in *ui-tests*.
 
 They make use of [playwright](https://playwright.dev/).
 
 The following commands run them:
 
 ```bash
-cd ./python/jupyterlab-chat/ui-tests
+cd ./ui-tests
 
 # Install the tests dependencies
 jlpm install
diff --git a/packages/jupyter-chat/package.json b/packages/jupyter-chat/package.json
index 860e2521..3ef7de7f 100644
--- a/packages/jupyter-chat/package.json
+++ b/packages/jupyter-chat/package.json
@@ -110,8 +110,7 @@
     "coverage",
     "**/*.d.ts",
     "tests",
-    "**/__tests__",
-    "ui-tests"
+    "**/__tests__"
   ],
   "eslintConfig": {
     "extends": [
diff --git a/packages/jupyterlab-chat/package.json b/packages/jupyterlab-chat/package.json
index cbfc6c49..d2ffab7e 100644
--- a/packages/jupyterlab-chat/package.json
+++ b/packages/jupyterlab-chat/package.json
@@ -114,8 +114,7 @@
     "coverage",
     "**/*.d.ts",
     "tests",
-    "**/__tests__",
-    "ui-tests"
+    "**/__tests__"
   ],
   "eslintConfig": {
     "extends": [
diff --git a/pyproject.toml b/pyproject.toml
index a465cbfd..29e0e9f3 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/README.md b/python/jupyterlab-chat/README.md
index e07e4568..851250a9 100644
--- a/python/jupyterlab-chat/README.md
+++ b/python/jupyterlab-chat/README.md
@@ -115,7 +115,7 @@ jlpm test
 This extension uses [Playwright](https://playwright.dev/docs/intro) for the integration tests (aka user level tests).
 More precisely, the JupyterLab helper [Galata](https://github.com/jupyterlab/jupyterlab/tree/master/galata) is used to handle testing the extension in JupyterLab.
 
-More information are provided within the [ui-tests](./ui-tests/README.md) README.
+More information are provided within the [ui-tests](../../ui-tests/README.md) README.
 
 ### Packaging the extension
 
diff --git a/python/jupyterlab-chat/jupyterlab_chat/py.typed b/python/jupyterlab-chat/jupyterlab_chat/py.typed
new file mode 100644
index 00000000..e69de29b
diff --git a/python/jupyterlab-chat/jupyterlab_chat/ychat.py b/python/jupyterlab-chat/jupyterlab_chat/ychat.py
index b3cdc172..519b53fa 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/package.json b/python/jupyterlab-chat/package.json
index 985a45a1..1a21e583 100644
--- a/python/jupyterlab-chat/package.json
+++ b/python/jupyterlab-chat/package.json
@@ -146,8 +146,7 @@
     "dist",
     "coverage",
     "**/*.d.ts",
-    "tests",
-    "ui-tests"
+    "tests"
   ],
   "eslintConfig": {
     "extends": [
diff --git a/python/jupyterlab-chat/pyproject.toml b/python/jupyterlab-chat/pyproject.toml
index 3d28553f..58146f93 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",
diff --git a/python/jupyterlab-chat/ui-tests/README.md b/ui-tests/README.md
similarity index 100%
rename from python/jupyterlab-chat/ui-tests/README.md
rename to ui-tests/README.md
diff --git a/python/jupyterlab-chat/ui-tests/jupyter_server_test_config.py b/ui-tests/jupyter_server_test_config.py
similarity index 100%
rename from python/jupyterlab-chat/ui-tests/jupyter_server_test_config.py
rename to ui-tests/jupyter_server_test_config.py
diff --git a/python/jupyterlab-chat/ui-tests/package.json b/ui-tests/package.json
similarity index 100%
rename from python/jupyterlab-chat/ui-tests/package.json
rename to ui-tests/package.json
diff --git a/python/jupyterlab-chat/ui-tests/playwright.config.js b/ui-tests/playwright.config.js
similarity index 100%
rename from python/jupyterlab-chat/ui-tests/playwright.config.js
rename to ui-tests/playwright.config.js
diff --git a/python/jupyterlab-chat/ui-tests/tests/autocompletion.spec.ts b/ui-tests/tests/autocompletion.spec.ts
similarity index 100%
rename from python/jupyterlab-chat/ui-tests/tests/autocompletion.spec.ts
rename to ui-tests/tests/autocompletion.spec.ts
diff --git a/python/jupyterlab-chat/ui-tests/tests/chat-file.spec.ts b/ui-tests/tests/chat-file.spec.ts
similarity index 100%
rename from python/jupyterlab-chat/ui-tests/tests/chat-file.spec.ts
rename to ui-tests/tests/chat-file.spec.ts
diff --git a/python/jupyterlab-chat/ui-tests/tests/code-toolbar.spec.ts b/ui-tests/tests/code-toolbar.spec.ts
similarity index 100%
rename from python/jupyterlab-chat/ui-tests/tests/code-toolbar.spec.ts
rename to ui-tests/tests/code-toolbar.spec.ts
diff --git a/python/jupyterlab-chat/ui-tests/tests/commands.spec.ts b/ui-tests/tests/commands.spec.ts
similarity index 100%
rename from python/jupyterlab-chat/ui-tests/tests/commands.spec.ts
rename to ui-tests/tests/commands.spec.ts
diff --git a/python/jupyterlab-chat/ui-tests/tests/commands.spec.ts-snapshots/launcher-tile-linux.png b/ui-tests/tests/commands.spec.ts-snapshots/launcher-tile-linux.png
similarity index 100%
rename from python/jupyterlab-chat/ui-tests/tests/commands.spec.ts-snapshots/launcher-tile-linux.png
rename to ui-tests/tests/commands.spec.ts-snapshots/launcher-tile-linux.png
diff --git a/python/jupyterlab-chat/ui-tests/tests/commands.spec.ts-snapshots/menu-new-linux.png b/ui-tests/tests/commands.spec.ts-snapshots/menu-new-linux.png
similarity index 100%
rename from python/jupyterlab-chat/ui-tests/tests/commands.spec.ts-snapshots/menu-new-linux.png
rename to ui-tests/tests/commands.spec.ts-snapshots/menu-new-linux.png
diff --git a/python/jupyterlab-chat/ui-tests/tests/message-toolbar.spec.ts b/ui-tests/tests/message-toolbar.spec.ts
similarity index 100%
rename from python/jupyterlab-chat/ui-tests/tests/message-toolbar.spec.ts
rename to ui-tests/tests/message-toolbar.spec.ts
diff --git a/python/jupyterlab-chat/ui-tests/tests/notifications.spec.ts b/ui-tests/tests/notifications.spec.ts
similarity index 100%
rename from python/jupyterlab-chat/ui-tests/tests/notifications.spec.ts
rename to ui-tests/tests/notifications.spec.ts
diff --git a/python/jupyterlab-chat/ui-tests/tests/notifications.spec.ts-snapshots/tab-with-unread-linux.png b/ui-tests/tests/notifications.spec.ts-snapshots/tab-with-unread-linux.png
similarity index 100%
rename from python/jupyterlab-chat/ui-tests/tests/notifications.spec.ts-snapshots/tab-with-unread-linux.png
rename to ui-tests/tests/notifications.spec.ts-snapshots/tab-with-unread-linux.png
diff --git a/python/jupyterlab-chat/ui-tests/tests/notifications.spec.ts-snapshots/tab-without-unread-linux.png b/ui-tests/tests/notifications.spec.ts-snapshots/tab-without-unread-linux.png
similarity index 100%
rename from python/jupyterlab-chat/ui-tests/tests/notifications.spec.ts-snapshots/tab-without-unread-linux.png
rename to ui-tests/tests/notifications.spec.ts-snapshots/tab-without-unread-linux.png
diff --git a/python/jupyterlab-chat/ui-tests/tests/raw-time.spec.ts b/ui-tests/tests/raw-time.spec.ts
similarity index 100%
rename from python/jupyterlab-chat/ui-tests/tests/raw-time.spec.ts
rename to ui-tests/tests/raw-time.spec.ts
diff --git a/python/jupyterlab-chat/ui-tests/tests/send-message.spec.ts b/ui-tests/tests/send-message.spec.ts
similarity index 100%
rename from python/jupyterlab-chat/ui-tests/tests/send-message.spec.ts
rename to ui-tests/tests/send-message.spec.ts
diff --git a/python/jupyterlab-chat/ui-tests/tests/side-panel.spec.ts b/ui-tests/tests/side-panel.spec.ts
similarity index 100%
rename from python/jupyterlab-chat/ui-tests/tests/side-panel.spec.ts
rename to ui-tests/tests/side-panel.spec.ts
diff --git a/python/jupyterlab-chat/ui-tests/tests/side-panel.spec.ts-snapshots/chat-icon-linux.png b/ui-tests/tests/side-panel.spec.ts-snapshots/chat-icon-linux.png
similarity index 100%
rename from python/jupyterlab-chat/ui-tests/tests/side-panel.spec.ts-snapshots/chat-icon-linux.png
rename to ui-tests/tests/side-panel.spec.ts-snapshots/chat-icon-linux.png
diff --git a/python/jupyterlab-chat/ui-tests/tests/side-panel.spec.ts-snapshots/moveToMain-linux.png b/ui-tests/tests/side-panel.spec.ts-snapshots/moveToMain-linux.png
similarity index 100%
rename from python/jupyterlab-chat/ui-tests/tests/side-panel.spec.ts-snapshots/moveToMain-linux.png
rename to ui-tests/tests/side-panel.spec.ts-snapshots/moveToMain-linux.png
diff --git a/python/jupyterlab-chat/ui-tests/tests/side-panel.spec.ts-snapshots/moveToSide-linux.png b/ui-tests/tests/side-panel.spec.ts-snapshots/moveToSide-linux.png
similarity index 100%
rename from python/jupyterlab-chat/ui-tests/tests/side-panel.spec.ts-snapshots/moveToSide-linux.png
rename to ui-tests/tests/side-panel.spec.ts-snapshots/moveToSide-linux.png
diff --git a/python/jupyterlab-chat/ui-tests/tests/test-utils.ts b/ui-tests/tests/test-utils.ts
similarity index 100%
rename from python/jupyterlab-chat/ui-tests/tests/test-utils.ts
rename to ui-tests/tests/test-utils.ts
diff --git a/python/jupyterlab-chat/ui-tests/tests/ui-config.spec.ts b/ui-tests/tests/ui-config.spec.ts
similarity index 100%
rename from python/jupyterlab-chat/ui-tests/tests/ui-config.spec.ts
rename to ui-tests/tests/ui-config.spec.ts
diff --git a/python/jupyterlab-chat/ui-tests/tests/ui-config.spec.ts-snapshots/not-stacked-messages-linux.png b/ui-tests/tests/ui-config.spec.ts-snapshots/not-stacked-messages-linux.png
similarity index 100%
rename from python/jupyterlab-chat/ui-tests/tests/ui-config.spec.ts-snapshots/not-stacked-messages-linux.png
rename to ui-tests/tests/ui-config.spec.ts-snapshots/not-stacked-messages-linux.png
diff --git a/python/jupyterlab-chat/ui-tests/tests/ui-config.spec.ts-snapshots/stacked-messages-linux.png b/ui-tests/tests/ui-config.spec.ts-snapshots/stacked-messages-linux.png
similarity index 100%
rename from python/jupyterlab-chat/ui-tests/tests/ui-config.spec.ts-snapshots/stacked-messages-linux.png
rename to ui-tests/tests/ui-config.spec.ts-snapshots/stacked-messages-linux.png
diff --git a/python/jupyterlab-chat/ui-tests/tests/unread.spec.ts b/ui-tests/tests/unread.spec.ts
similarity index 100%
rename from python/jupyterlab-chat/ui-tests/tests/unread.spec.ts
rename to ui-tests/tests/unread.spec.ts
diff --git a/python/jupyterlab-chat/ui-tests/tests/unread.spec.ts-snapshots/navigation-bottom-linux.png b/ui-tests/tests/unread.spec.ts-snapshots/navigation-bottom-linux.png
similarity index 100%
rename from python/jupyterlab-chat/ui-tests/tests/unread.spec.ts-snapshots/navigation-bottom-linux.png
rename to ui-tests/tests/unread.spec.ts-snapshots/navigation-bottom-linux.png
diff --git a/python/jupyterlab-chat/ui-tests/tests/unread.spec.ts-snapshots/navigation-bottom-unread-linux.png b/ui-tests/tests/unread.spec.ts-snapshots/navigation-bottom-unread-linux.png
similarity index 100%
rename from python/jupyterlab-chat/ui-tests/tests/unread.spec.ts-snapshots/navigation-bottom-unread-linux.png
rename to ui-tests/tests/unread.spec.ts-snapshots/navigation-bottom-unread-linux.png
diff --git a/python/jupyterlab-chat/ui-tests/tests/unread.spec.ts-snapshots/navigation-top-linux.png b/ui-tests/tests/unread.spec.ts-snapshots/navigation-top-linux.png
similarity index 100%
rename from python/jupyterlab-chat/ui-tests/tests/unread.spec.ts-snapshots/navigation-top-linux.png
rename to ui-tests/tests/unread.spec.ts-snapshots/navigation-top-linux.png
diff --git a/python/jupyterlab-chat/ui-tests/yarn.lock b/ui-tests/yarn.lock
similarity index 100%
rename from python/jupyterlab-chat/ui-tests/yarn.lock
rename to ui-tests/yarn.lock