diff --git a/.github/workflows/unit-tests.yml b/.github/workflows/unit-tests.yml
index 7ece99976..05224180e 100644
--- a/.github/workflows/unit-tests.yml
+++ b/.github/workflows/unit-tests.yml
@@ -1,4 +1,4 @@
-name: Python Unit Tests
+name: Python Tests
 
 # suppress warning raised by https://github.com/jupyter/jupyter_core/pull/292
 env:
@@ -12,7 +12,7 @@ on:
 
 jobs:
   unit-tests:
-    name: Linux
+    name: Unit tests
     runs-on: ubuntu-latest
     steps:
       - name: Checkout
@@ -30,7 +30,7 @@ jobs:
           pytest -vv -r ap --cov jupyter_ai
 
   typing-tests:
-    name: Linux
+    name: Typing test
     runs-on: ubuntu-latest
     steps:
       - name: Checkout
@@ -42,7 +42,7 @@ jobs:
       - name: Install extension dependencies and build the extension
         run: ./scripts/install.sh
 
-      - name: Execute unit tests
+      - name: Run mypy
         run: |
           set -eux
           mypy --version
diff --git a/packages/jupyter-ai/pyproject.toml b/packages/jupyter-ai/pyproject.toml
index e8deeb133..093b2159c 100644
--- a/packages/jupyter-ai/pyproject.toml
+++ b/packages/jupyter-ai/pyproject.toml
@@ -50,6 +50,8 @@ test = [
     "pytest-tornasync",
     "pytest-jupyter",
     "syrupy~=4.0.8",
+    "types-jsonschema",
+    "mypy"
 ]
 
 dev = ["jupyter_ai_magics[dev]"]