reflex 0.8.15a1__py3-none-any.whl → 0.8.16a1__py3-none-any.whl

reflex/istate/manager/disk.py

@@ -4,15 +4,27 @@ import asyncio
 import contextlib
 import dataclasses
 import functools
+import time
 from collections.abc import AsyncIterator
 from hashlib import md5
 from pathlib import Path
 
 from typing_extensions import override
 
+from reflex.environment import environment
 from reflex.istate.manager import StateManager, _default_token_expiration
 from reflex.state import BaseState, _split_substate_key, _substate_key
-from reflex.utils import path_ops, prerequisites
+from reflex.utils import console, path_ops, prerequisites
+from reflex.utils.misc import run_in_thread
+
+
+@dataclasses.dataclass(frozen=True)
+class QueueItem:
+    """An item in the write queue."""
+
+    token: str
+    state: BaseState
+    timestamp: float
 
 
 @dataclasses.dataclass
@@ -34,6 +46,22 @@ class StateManagerDisk(StateManager):
     # The token expiration time (s).
     token_expiration: int = dataclasses.field(default_factory=_default_token_expiration)
 
+    # Last time a token was touched.
+    _token_last_touched: dict[str, float] = dataclasses.field(
+        default_factory=dict,
+        init=False,
+    )
+
+    # Pending writes
+    _write_queue: dict[str, QueueItem] = dataclasses.field(
+        default_factory=dict,
+        init=False,
+    )
+    _write_queue_task: asyncio.Task | None = None
+    _write_debounce_seconds: float = dataclasses.field(
+        default=environment.REFLEX_STATE_MANAGER_DISK_DEBOUNCE_SECONDS.get()
+    )
+
     def __post_init__(self):
         """Create a new state manager."""
         path_ops.mkdir(self.states_directory)
@@ -51,8 +79,6 @@ class StateManagerDisk(StateManager):
 
     def _purge_expired_states(self):
         """Purge expired states from the disk."""
-        import time
-
         for path in path_ops.ls(self.states_directory):
             # check path is a pickle file
             if path.suffix != ".pkl":
@@ -137,6 +163,7 @@ class StateManagerDisk(StateManager):
             The state for the token.
         """
         client_token = _split_substate_key(token)[0]
+        self._token_last_touched[client_token] = time.time()
         root_state = self.states.get(client_token)
         if root_state is not None:
             # Retrieved state from memory.
@@ -170,11 +197,109 @@ class StateManagerDisk(StateManager):
             if pickle_state:
                 if not self.states_directory.exists():
                     self.states_directory.mkdir(parents=True, exist_ok=True)
-                self.token_path(substate_token).write_bytes(pickle_state)
+                await run_in_thread(
+                    lambda: self.token_path(substate_token).write_bytes(pickle_state),
+                )
 
         for substate_substate in substate.substates.values():
             await self.set_state_for_substate(client_token, substate_substate)
 
+    async def _process_write_queue_delay(self):
+        """Wait for the debounce period before processing the write queue again."""
+        now = time.time()
+        if self._write_queue:
+            # There are still items in the queue, schedule another run.
+            next_write_in = max(
+                0,
+                min(
+                    self._write_debounce_seconds - (now - item.timestamp)
+                    for item in self._write_queue.values()
+                ),
+            )
+            await asyncio.sleep(next_write_in)
+        elif self._write_debounce_seconds > 0:
+            # No items left, wait a bit before checking again.
+            await asyncio.sleep(self._write_debounce_seconds)
+        else:
+            # Debounce is disabled, so sleep until the next token expiration.
+            oldest_token_last_touch = min(
+                self._token_last_touched.values(), default=now
+            )
+            next_expiration_in = self.token_expiration - (now - oldest_token_last_touch)
+            await asyncio.sleep(next_expiration_in)
+
+    async def _process_write_queue(self):
+        """Long running task that checks for states to write to disk.
+
+        Raises:
+            asyncio.CancelledError: When the task is cancelled.
+        """
+        while True:
+            try:
+                now = time.time()
+                # sort the _write_queue by oldest timestamp and exclude items younger than debounce time
+                items_to_write = sorted(
+                    (
+                        item
+                        for item in self._write_queue.values()
+                        if now - item.timestamp >= self._write_debounce_seconds
+                    ),
+                    key=lambda item: item.timestamp,
+                )
+                for item in items_to_write:
+                    token = item.token
+                    client_token, _ = _split_substate_key(token)
+                    await self.set_state_for_substate(
+                        client_token, self._write_queue.pop(token).state
+                    )
+                # Check for expired states to purge.
+                for token, last_touched in list(self._token_last_touched.items()):
+                    if now - last_touched > self.token_expiration:
+                        self._token_last_touched.pop(token)
+                        self.states.pop(token, None)
+                await run_in_thread(self._purge_expired_states)
+                await self._process_write_queue_delay()
+            except asyncio.CancelledError:  # noqa: PERF203
+                await self._flush_write_queue()
+                raise
+            except Exception as e:
+                console.error(f"Error processing write queue: {e!r}")
+                if e.args == ("cannot schedule new futures after shutdown",):
+                    # Event loop is shutdown, nothing else we can really do...
+                    return
+                await self._process_write_queue_delay()
+
+    async def _flush_write_queue(self):
+        """Flush any remaining items in the write queue to disk."""
+        outstanding_items = list(self._write_queue.values())
+        n_outstanding_items = len(outstanding_items)
+        self._write_queue.clear()
+        # When the task is cancelled, write all remaining items to disk.
+        console.debug(
+            f"StateManagerDisk._flush_write_queue: writing {n_outstanding_items} remaining items to disk"
+        )
+        for item in outstanding_items:
+            token = item.token
+            client_token, _ = _split_substate_key(token)
+            await self.set_state_for_substate(
+                client_token,
+                item.state,
+            )
+        console.debug(
+            f"StateManagerDisk._flush_write_queue: Finished writing {n_outstanding_items} items"
+        )
+
+    async def _schedule_process_write_queue(self):
+        """Schedule the write queue processing task if not already running."""
+        if self._write_queue_task is None or self._write_queue_task.done():
+            async with self._state_manager_lock:
+                if self._write_queue_task is None or self._write_queue_task.done():
+                    self._write_queue_task = asyncio.create_task(
+                        self._process_write_queue(),
+                        name="StateManagerDisk|WriteQueueProcessor",
+                    )
+                    await asyncio.sleep(0)  # Yield to allow the task to start.
+
     @override
     async def set_state(self, token: str, state: BaseState):
         """Set the state for a token.
@@ -184,7 +309,19 @@ class StateManagerDisk(StateManager):
             state: The state to set.
         """
         client_token, _ = _split_substate_key(token)
-        await self.set_state_for_substate(client_token, state)
+        if self._write_debounce_seconds > 0:
+            # Deferred write to reduce disk IO overhead.
+            if client_token not in self._write_queue:
+                self._write_queue[client_token] = QueueItem(
+                    token=client_token,
+                    state=state,
+                    timestamp=time.time(),
+                )
+        else:
+            # Immediate write to disk.
+            await self.set_state_for_substate(client_token, state)
+        # Ensure the processing task is scheduled to handle expirations and any deferred writes.
+        await self._schedule_process_write_queue()
 
     @override
     @contextlib.asynccontextmanager
@@ -208,3 +345,12 @@ class StateManagerDisk(StateManager):
             state = await self.get_state(token)
             yield state
             await self.set_state(token, state)
+
+    async def close(self):
+        """Close the state manager, flushing any pending writes to disk."""
+        async with self._state_manager_lock:
+            if self._write_queue_task:
+                self._write_queue_task.cancel()
+                with contextlib.suppress(asyncio.CancelledError):
+                    await self._write_queue_task
+                    self._write_queue_task = None

reflex/model.py

@@ -299,7 +299,7 @@ if find_spec("sqlmodel") and find_spec("sqlalchemy") and find_spec("pydantic"):
         # Format output with message
         return f"  [{status_icon}] {current}{head_marker}, {message}"
 
-    async def get_db_status() -> dict[str, bool]:
+    def get_db_status() -> dict[str, bool]:
         """Checks the status of the database connection.
 
         Attempts to connect to the database and execute a simple query to verify connectivity.

reflex/plugins/_screenshot.py

@@ -68,7 +68,7 @@ class ScreenshotPlugin(BasePlugin):
         if not app._api:
             return
 
-        async def active_connections(_request: "Request") -> "Response":
+        def active_connections(_request: "Request") -> "Response":
             from starlette.responses import JSONResponse
 
             if not app.event_namespace:
@@ -122,7 +122,7 @@ class ScreenshotPlugin(BasePlugin):
             while found_new:
                 found_new = False
 
-                for state in all_states:
+                for state in list(all_states):
                     for substate in state.substates.values():
                         substate._was_touched = True
 

reflex/plugins/shared_tailwind.py

@@ -112,12 +112,9 @@ def tailwind_config_js_template(
     ]
 
     # Generate import statements for destructured imports
-    import_lines = "\n".join(
-        [
-            f"import {{ {imp['name']} }} from {json.dumps(imp['from'])};"
-            for imp in imports
-        ]
-    )
+    import_lines = "\n".join([
+        f"import {{ {imp['name']} }} from {json.dumps(imp['from'])};" for imp in imports
+    ])
 
     # Generate plugin imports
     plugin_imports = []
@@ -131,12 +128,10 @@ def tailwind_config_js_template(
 
     plugin_imports_lines = "\n".join(plugin_imports)
 
-    presets_imports_lines = "\n".join(
-        [
-            f"import preset{i} from {json.dumps(preset)};"
-            for i, preset in enumerate(presets, 1)
-        ]
-    )
+    presets_imports_lines = "\n".join([
+        f"import preset{i} from {json.dumps(preset)};"
+        for i, preset in enumerate(presets, 1)
+    ])
 
     # Generate plugin array
     plugin_list = []
@@ -159,8 +154,8 @@ def tailwind_config_js_template(
 {presets_imports_lines}
 
 export default {{
-    content: {json.dumps(content if content else default_content)},
-    theme: {json.dumps(theme if theme else {})},
+    content: {json.dumps(content or default_content)},
+    theme: {json.dumps(theme or {})},
     {f"darkMode: {json.dumps(dark_mode)}," if dark_mode is not None else ""}
     {f"corePlugins: {json.dumps(core_plugins)}," if core_plugins is not None else ""}
     {f"importants: {json.dumps(important)}," if important is not None else ""}

reflex/reflex.py

@@ -269,15 +269,13 @@ def _run(
 
     # In prod mode, run the backend on a separate thread.
     if backend and env == constants.Env.PROD:
-        commands.append(
-            (
-                backend_cmd,
-                backend_host,
-                backend_port,
-                config.loglevel.subprocess_level(),
-                frontend,
-            )
-        )
+        commands.append((
+            backend_cmd,
+            backend_host,
+            backend_port,
+            config.loglevel.subprocess_level(),
+            frontend,
+        ))
 
     if single_port:
         setup_frontend(Path.cwd())

reflex/state.py

@@ -112,13 +112,13 @@ def _no_chain_background_task(state: BaseState, name: str, fn: Callable) -> Call
     )
     if inspect.iscoroutinefunction(fn):
 
-        async def _no_chain_background_task_co(*args, **kwargs):
+        async def _no_chain_background_task_co(*args, **kwargs):  # noqa: RUF029
             raise RuntimeError(message)
 
         return _no_chain_background_task_co
     if inspect.isasyncgenfunction(fn):
 
-        async def _no_chain_background_task_gen(*args, **kwargs):
+        async def _no_chain_background_task_gen(*args, **kwargs):  # noqa: RUF029
             yield
             raise RuntimeError(message)
 
@@ -187,14 +187,12 @@ class EventHandlerSetVar(EventHandler):
         Returns:
             The hash of the event handler.
         """
-        return hash(
-            (
-                tuple(self.event_actions.items()),
-                self.fn,
-                self.state_full_name,
-                self.state_cls,
-            )
-        )
+        return hash((
+            tuple(self.event_actions.items()),
+            self.fn,
+            self.state_full_name,
+            self.state_cls,
+        ))
 
     def setvar(self, var_name: str, value: Any):
         """Set the state variable to the value of the event.
@@ -529,14 +527,12 @@ class BaseState(EvenMoreBasicBaseState):
             if types.is_backend_base_variable(name, cls)
         }
         # Add annotated backend vars that may not have a default value.
-        new_backend_vars.update(
-            {
-                name: cls._get_var_default(name, annotation_value)
-                for name, annotation_value in cls._get_type_hints().items()
-                if name not in new_backend_vars
-                and types.is_backend_base_variable(name, cls)
-            }
-        )
+        new_backend_vars.update({
+            name: cls._get_var_default(name, annotation_value)
+            for name, annotation_value in cls._get_type_hints().items()
+            if name not in new_backend_vars
+            and types.is_backend_base_variable(name, cls)
+        })
 
         cls.backend_vars = {
             **cls.inherited_backend_vars,
@@ -794,9 +790,10 @@ class BaseState(EvenMoreBasicBaseState):
                         parent_state = defining_state_cls.get_parent_state()
                         if parent_state is not None:
                             defining_state_cls = parent_state
-                    defining_state_cls._var_dependencies.setdefault(dvar, set()).add(
-                        (cls.get_full_name(), cvar_name)
-                    )
+                    defining_state_cls._var_dependencies.setdefault(dvar, set()).add((
+                        cls.get_full_name(),
+                        cvar_name,
+                    ))
                     defining_state_cls._potentially_dirty_states.add(
                         cls.get_full_name()
                     )
@@ -1143,7 +1140,7 @@ class BaseState(EvenMoreBasicBaseState):
         from reflex.config import get_config
 
         config = get_config()
-        _create_event_handler_kwargs = {}
+        create_event_handler_kwargs = {}
 
         if config.state_auto_setters is None:
 
@@ -1159,14 +1156,14 @@ class BaseState(EvenMoreBasicBaseState):
                     )
                     return super().__call__(*args, **kwargs)
 
-            _create_event_handler_kwargs["event_handler_cls"] = (
+            create_event_handler_kwargs["event_handler_cls"] = (
                 EventHandlerDeprecatedSetter
             )
 
         setter_name = Var._get_setter_name_for_name(name)
         if setter_name not in cls.__dict__:
             event_handler = cls._create_event_handler(
-                prop._get_setter(name), **_create_event_handler_kwargs
+                prop._get_setter(name), **create_event_handler_kwargs
             )
             cls.event_handlers[setter_name] = event_handler
             setattr(cls, setter_name, event_handler)
@@ -1524,12 +1521,10 @@ class BaseState(EvenMoreBasicBaseState):
         return {
             cls.get_class_substate(substate_name)
             for substate_name in cls._always_dirty_substates
-        }.union(
-            {
-                cls.get_root_state().get_class_substate(substate_name)
-                for substate_name in cls._potentially_dirty_states
-            }
-        )
+        }.union({
+            cls.get_root_state().get_class_substate(substate_name)
+            for substate_name in cls._potentially_dirty_states
+        })
 
     def _get_root_state(self) -> BaseState:
         """Get the root state of the state tree.
@@ -1858,13 +1853,11 @@ class BaseState(EvenMoreBasicBaseState):
             ):
                 if issubclass(hinted_args, Model):
                     # Remove non-fields from the payload
-                    payload[arg] = hinted_args(
-                        **{
-                            key: value
-                            for key, value in value.items()
-                            if key in hinted_args.__fields__
-                        }
-                    )
+                    payload[arg] = hinted_args(**{
+                        key: value
+                        for key, value in value.items()
+                        if key in hinted_args.__fields__
+                    })
                 elif dataclasses.is_dataclass(hinted_args):
                     payload[arg] = hinted_args(**value)
                 elif find_spec("pydantic"):

reflex/style.py

@@ -289,11 +289,11 @@ class Style(dict[str, Any]):
             value: The value to set.
         """
         # Create a Var to collapse VarData encoded in f-string.
-        _var = LiteralVar.create(value)
-        if _var is not None:
+        var = LiteralVar.create(value)
+        if var is not None:
             # Carry the imports/hooks when setting a Var as a value.
             self._var_data = VarData.merge(
-                getattr(self, "_var_data", None), _var._get_all_var_data()
+                getattr(self, "_var_data", None), var._get_all_var_data()
             )
         super().__setitem__(key, value)
 
@@ -348,7 +348,7 @@ def format_as_emotion(style_dict: dict[str, Any]) -> Style | None:
     Returns:
         The emotion style dict.
     """
-    _var_data = style_dict._var_data if isinstance(style_dict, Style) else None
+    var_data = style_dict._var_data if isinstance(style_dict, Style) else None
 
     emotion_style = Style()
 
@@ -381,8 +381,8 @@ def format_as_emotion(style_dict: dict[str, Any]) -> Style | None:
         else:
             emotion_style[key] = value
     if emotion_style:
-        if _var_data is not None:
-            emotion_style._var_data = VarData.merge(emotion_style._var_data, _var_data)
+        if var_data is not None:
+            emotion_style._var_data = VarData.merge(emotion_style._var_data, var_data)
         return emotion_style
     return None
 

reflex/testing.py

@@ -28,10 +28,8 @@ from typing import TYPE_CHECKING, Any, Literal, TypeVar
 import uvicorn
 
 import reflex
-import reflex.environment
 import reflex.reflex
 import reflex.utils.build
-import reflex.utils.exec
 import reflex.utils.format
 import reflex.utils.prerequisites
 import reflex.utils.processes
@@ -41,7 +39,12 @@ from reflex.environment import environment
 from reflex.istate.manager.disk import StateManagerDisk
 from reflex.istate.manager.memory import StateManagerMemory
 from reflex.istate.manager.redis import StateManagerRedis
-from reflex.state import BaseState, StateManager, reload_state_module
+from reflex.state import (
+    BaseState,
+    StateManager,
+    _split_substate_key,
+    reload_state_module,
+)
 from reflex.utils import console, js_runtimes
 from reflex.utils.export import export
 from reflex.utils.token_manager import TokenManager
@@ -246,14 +249,12 @@ class AppHarness:
             if isinstance(self.app_source, functools.partial):
                 self.app_source = self.app_source.func
             # get the source from a function or module object
-            source_code = "\n".join(
-                [
-                    "\n".join(
-                        self.get_app_global_source(k, v) for k, v in app_globals.items()
-                    ),
-                    self._get_source_from_app_source(self.app_source),
-                ]
-            )
+            source_code = "\n".join([
+                "\n".join([
+                    self.get_app_global_source(k, v) for k, v in app_globals.items()
+                ]),
+                self._get_source_from_app_source(self.app_source),
+            ])
             get_config().loglevel = reflex.constants.LogLevel.INFO
             with chdir(self.app_path):
                 reflex.reflex._init(
@@ -279,15 +280,16 @@ class AppHarness:
                 )
             )
             self.app_asgi = self.app_instance()
-        if self.app_instance and isinstance(
-            self.app_instance._state_manager, StateManagerRedis
-        ):
+        if self.app_instance and self.app_instance._state_manager is not None:
             if self.app_instance._state is None:
                 msg = "State is not set."
                 raise RuntimeError(msg)
-            # Create our own redis connection for testing.
-            self.state_manager = StateManagerRedis.create(self.app_instance._state)
-        else:
+            if isinstance(self.app_instance._state_manager, StateManagerRedis):
+                # Create our own redis connection for testing.
+                self.state_manager = StateManagerRedis.create(self.app_instance._state)
+            elif isinstance(self.app_instance._state_manager, StateManagerDisk):
+                self.state_manager = StateManagerDisk.create(self.app_instance._state)
+        if self.state_manager is None:
             self.state_manager = (
                 self.app_instance._state_manager if self.app_instance else None
             )
@@ -305,8 +307,9 @@ class AppHarness:
 
         async def _shutdown(*args, **kwargs) -> None:
             # ensure redis is closed before event loop
-            if self.app_instance is not None and isinstance(
-                self.app_instance._state_manager, StateManagerRedis
+            if (
+                self.app_instance is not None
+                and self.app_instance._state_manager is not None
             ):
                 with contextlib.suppress(ValueError):
                     await self.app_instance._state_manager.close()
@@ -358,6 +361,12 @@ class AppHarness:
         Raises:
             RuntimeError: when the state manager cannot be reset
         """
+        if (
+            self.app_instance is not None
+            and self.app_instance._state_manager is not None
+        ):
+            with contextlib.suppress(RuntimeError):
+                await self.app_instance._state_manager.close()
         if (
             self.app_instance is not None
             and isinstance(
@@ -366,8 +375,6 @@ class AppHarness:
             )
             and self.app_instance._state is not None
         ):
-            with contextlib.suppress(RuntimeError):
-                await self.app_instance._state_manager.close()
             self.app_instance._state_manager = StateManagerRedis.create(
                 state=self.app_instance._state,
             )
@@ -713,11 +720,20 @@ class AppHarness:
         if self.state_manager is None:
             msg = "state_manager is not set."
             raise RuntimeError(msg)
+        if self.app_instance is not None and isinstance(
+            self.app_instance.state_manager, StateManagerDisk
+        ):
+            # Song and dance to convince the instance's state manager to flush
+            # (we can't directly await the _other_ loop's Future)
+            await self.app_instance.state_manager._flush_write_queue()
+        if isinstance(self.state_manager, StateManagerDisk):
+            # Force reload the latest state from disk.
+            client_token, _ = _split_substate_key(token)
+            self.state_manager.states.pop(client_token, None)
         try:
             return await self.state_manager.get_state(token)
         finally:
-            if isinstance(self.state_manager, StateManagerRedis):
-                await self.state_manager.close()
+            await self.state_manager.close()
 
     async def set_state(self, token: str, **kwargs) -> None:
         """Set the state associated with the given token.
@@ -738,8 +754,13 @@ class AppHarness:
         try:
             await self.state_manager.set_state(token, state)
         finally:
-            if isinstance(self.state_manager, StateManagerRedis):
-                await self.state_manager.close()
+            if self.app_instance is not None and isinstance(
+                self.app_instance.state_manager, StateManagerDisk
+            ):
+                # Clear the token from the backend's cache so it will be reloaded.
+                client_token, _ = _split_substate_key(token)
+                self.app_instance.state_manager.states.pop(client_token, None)
+            await self.state_manager.close()
 
     @contextlib.asynccontextmanager
     async def modify_state(self, token: str) -> AsyncIterator[BaseState]:
@@ -761,17 +782,20 @@ class AppHarness:
             msg = "App is not running."
             raise RuntimeError(msg)
         app_state_manager = self.app_instance.state_manager
-        if isinstance(self.state_manager, StateManagerRedis):
+        if isinstance(self.state_manager, (StateManagerRedis, StateManagerDisk)):
             # Temporarily replace the app's state manager with our own, since
-            # the redis connection is on the backend_thread event loop
+            # the redis/disk connection is on the backend_thread event loop
             self.app_instance._state_manager = self.state_manager
         try:
             async with self.app_instance.modify_state(token) as state:
                 yield state
         finally:
-            if isinstance(self.state_manager, StateManagerRedis):
-                self.app_instance._state_manager = app_state_manager
-                await self.state_manager.close()
+            if isinstance(app_state_manager, StateManagerDisk):
+                # Clear the token from the cache so it will be reloaded.
+                client_token, _ = _split_substate_key(token)
+                app_state_manager.states.pop(client_token, None)
+            await self.state_manager.close()
+            self.app_instance._state_manager = app_state_manager
 
     def token_manager(self) -> TokenManager:
         """Get the token manager for the app instance.