diff --git a/tests/test_engine.py b/tests/test_engine.py
index 8a5ec80..1c19c56 100644
--- a/tests/test_engine.py
+++ b/tests/test_engine.py
@@ -471,13 +471,12 @@ async def task_b() -> str:
     tasks_dict_ref = None
 
     # We need to capture the reference inside the execution scope
-    original_execute = engine.execute
 
     async def mock_execute() -> dict[str, typing.Any]:
         nonlocal tasks_dict_ref
         engine._cached_topo_order = ["A", "B"]
 
-        class TasksDict(dict):
+        class TasksDict(dict[str, asyncio.Task[typing.Any]]):
             pass
         tasks = TasksDict()
         tasks_dict_ref = weakref.ref(tasks)
@@ -501,7 +500,7 @@ class TasksDict(dict):
 
     gc.collect()
 
-    assert tasks_dict_ref() is None, (
+    assert tasks_dict_ref is not None and tasks_dict_ref() is None, (
         "Memory leak: tasks dictionary was not garbage collected"
     )
 
diff --git a/tests/test_memory_leak.py b/tests/test_memory_leak.py
index d1babdf..ce5e313 100644
--- a/tests/test_memory_leak.py
+++ b/tests/test_memory_leak.py
@@ -1,8 +1,9 @@
+import typing
 import asyncio
 import gc
 import weakref
 import pytest
-from catalyst.domain.engine import WorkflowEngine, TaskError
+from catalyst.domain.engine import WorkflowEngine
 
 
 @pytest.mark.asyncio
@@ -43,7 +44,7 @@ async def test_reference_cycle_is_broken() -> None:
     """Ensure that the execution dictionary is not held in a reference cycle."""
     engine = WorkflowEngine()
 
-    async def my_task():
+    async def my_task() -> str:
         return "success"
 
     engine.add_task("task_a", my_task)
@@ -52,22 +53,25 @@ async def my_task():
     results = await engine.execute()
     assert results["task_a"] == "success"
 
-    weak_dep = None
+    weak_dep: typing.Any = None
 
     orig_run_node = engine._run_node
 
     # We create a dummy class to hold a reference to tasks so we can weakref it
     class TaskHolder:
-        def __init__(self, tasks):
+        def __init__(self, tasks: tuple[asyncio.Task[typing.Any], ...]) -> None:
             self.tasks = tasks
 
-    async def wrapped_run_node(node: str, dep_tasks: tuple[asyncio.Task, ...]):
+
+
+
+    async def wrapped_run_node(node: str, dep_tasks: tuple[asyncio.Task[typing.Any], ...]) -> typing.Any:
         nonlocal weak_dep
         holder = TaskHolder(dep_tasks)
         weak_dep = weakref.ref(holder)
         return await orig_run_node(node, dep_tasks)
 
-    engine._run_node = wrapped_run_node
+    engine._run_node = wrapped_run_node  # type: ignore
 
     await engine.execute()
     gc.collect()
diff --git a/uv.lock b/uv.lock
index eb4ba1f..b38eb30 100644
--- a/uv.lock
+++ b/uv.lock
@@ -49,7 +49,7 @@ wheels = [
 
 [[package]]
 name = "catalyst"
-version = "0.1.29"
+version = "0.1.31"
 source = { editable = "." }
 dependencies = [
     { name = "fastapi" },