diff --git a/.github/workflows/unit-tests.yml b/.github/workflows/unit-tests.yml
index 7ece99976..3f6a275fa 100644
--- a/.github/workflows/unit-tests.yml
+++ b/.github/workflows/unit-tests.yml
@@ -11,8 +11,8 @@ on:
     branches: "*"
 
 jobs:
-  unit-tests:
-    name: Linux
+  setup:
+    name: Setup
     runs-on: ubuntu-latest
     steps:
       - name: Checkout
@@ -24,25 +24,22 @@ jobs:
       - name: Install extension dependencies and build the extension
         run: ./scripts/install.sh
 
+  unit-tests:
+    name: Unit tests
+    needs: setup
+    runs-on: ubuntu-latest
+    steps:
       - name: Execute unit tests
         run: |
           set -eux
           pytest -vv -r ap --cov jupyter_ai
 
   typing-tests:
-    name: Linux
+    name: Typing test
+    needs: setup
     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/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]"]