diff --git a/justfile b/justfile
index d6deb6f8..8718b6b9 100644
--- a/justfile
+++ b/justfile
@@ -1,3 +1,3 @@
 regenerate-protobufs:
-    protoc --proto_path=shared/protobufs/schemas --python_out=shared/protobufs/types shared/protobufs/schemas/*.proto
+    protoc --proto_path=shared/protobufs/schemas --python_out=shared/protobufs/types --pyi_out=shared/protobufs/types shared/protobufs/schemas/*.proto
     uv run ruff format ./shared/protobufs/types
\ No newline at end of file
diff --git a/pyproject.toml b/pyproject.toml
index b5e485c8..baea585c 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -29,7 +29,7 @@ darwin = [
 fmt = { shell = "ruff format master worker shared engines/*", help = "Format the code" }
 lint = { shell = "ruff check --fix master worker shared engines/*", help = "Run the linter" }
 test = { shell = "pytest master worker shared engines/*", help = "Run the tests" }
-check = { shell = "basedpyright --project master worker shared engines/*", help = "Run type checker" }
+check = { shell = "basedpyright --project pyproject.toml", help = "Run type checker" }
 sync = { shell = "uv sync --all-packages", help = "Sync the dependencies" }
 protobufs = { shell = "just regenerate-protobufs", help = "Regenerate the protobufs" }
 build = { shell = "just regenerate-protobufs && uv build --all-packages", help = "Build the project" }
@@ -85,7 +85,10 @@ include = ["master", "worker", "shared", "engines/*"]
 pythonVersion  = "3.13"
 pythonPlatform = "Darwin"
 
-exclude = ["shared/protobufs/**"]
+stubPath = "shared/protobufs/types"
+ignore   = [
+    "shared/protobufs/types/**/*",
+]
 
 ###
 # uv configuration
diff --git a/uv.lock b/uv.lock
index 5f8ae494..dca182fb 100644
--- a/uv.lock
+++ b/uv.lock
@@ -64,7 +64,7 @@ dev = [
 
 [package.metadata]
 requires-dist = [
-    { name = "exo-master", virtual = "master" },
+    { name = "exo-master", editable = "master" },
     { name = "exo-worker", virtual = "worker" },
     { name = "mlx", marker = "extra == 'darwin'" },
 ]
@@ -81,12 +81,12 @@ dev = [
 [[package]]
 name = "exo-engine-mlx"
 version = "0.1.0"
-source = { virtual = "engines/mlx" }
+source = { editable = "engines/mlx" }
 
 [[package]]
 name = "exo-master"
 version = "0.1.0"
-source = { virtual = "master" }
+source = { editable = "master" }
 
 [[package]]
 name = "exo-shared"