diff --git a/reflex/istate/proxy.py b/reflex/istate/proxy.py
index ad97d59162..96abb22af3 100644
--- a/reflex/istate/proxy.py
+++ b/reflex/istate/proxy.py
@@ -14,6 +14,7 @@
 from typing import TYPE_CHECKING, Any, SupportsIndex, TypeVar
 
 import wrapt
+from typing_extensions import Self
 
 from reflex.base import Base
 from reflex.utils import prerequisites
@@ -93,7 +94,7 @@ def _is_mutable(self) -> bool:
             return self._self_parent_state_proxy._is_mutable() or self._self_mutable
         return self._self_mutable
 
-    async def __aenter__(self) -> StateProxy:
+    async def __aenter__(self) -> Self:
         """Enter the async context manager protocol.
 
         Sets mutability to True and enters the `App.modify_state` async context,
diff --git a/reflex/state.py b/reflex/state.py
index 8ffe0abf3b..daf33e2631 100644
--- a/reflex/state.py
+++ b/reflex/state.py
@@ -2236,7 +2236,7 @@ def dict(
 
         return d
 
-    async def __aenter__(self) -> BaseState:
+    async def __aenter__(self) -> Self:
         """Enter the async context manager protocol.
 
         This is a no-op for the State class and mainly used in background-tasks/StateProxy.