add config option to allow state mutation with contextmanager from non-background events

Author: benedikt-bartscher

reflex/environment.py

@@ -727,6 +727,11 @@ class EnvironmentVariables:
     # How long to wait between automatic reload on frontend error to avoid reload loops.
     REFLEX_AUTO_RELOAD_COOLDOWN_TIME_MS: EnvVar[int] = env_var(10_000)
 
+    # Wheter to allow entering a state mutation with `async with self` in normal (non background) event handlers
+    REFLEX_STATE_ALLOW_CONTEXTMANAGER_WITHOUT_BACKGROUND_TASK: EnvVar[bool] = env_var(
+        False
+    )
+
 
 environment = EnvironmentVariables()
 

reflex/state.py

@@ -2192,9 +2192,14 @@ async def __aenter__(self) -> BaseState:
         This should not be used for the State class, but exists for
         type-compatibility with StateProxy.
 
+        Returns:
+            The unmodified state (self)
+
         Raises:
             TypeError: always, because async contextmanager protocol is only supported for background task.
         """
+        if environment.REFLEX_STATE_ALLOW_CONTEXTMANAGER_WITHOUT_BACKGROUND_TASK.get():
+            return self
         msg = "Only background task should use `async with self` to modify state."
         raise TypeError(msg)