laipeng/reflex
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Cache root_state under client_token for StateManagerMemory compatibility

Browse Source 
Mainly this just makes it easier for us to write tests that work against either
Disk or Memory state managers.
This commit is contained in:
Masen Furer 2024-10-03 19:40:46 -07:00
parent 20d8192360
commit 16db3d3905
No known key found for this signature in database
GPG Key ID: B0008AD22B3B3A95
1 changed files with 3 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
reflex/state.py
View File

@ -2764,14 +2764,13 @@ class StateManagerDisk(StateManager):
The state for the token. The state for the token.
""" """
client_token = _split_substate_key(token)[0] client_token = _split_substate_key(token)[0]
root_state_token = _substate_key(client_token, self.state) root_state = self.states.get(client_token)
root_state = self.states.get(root_state_token)
if root_state is not None: if root_state is not None:
# Retrieved state from memory. # Retrieved state from memory.
return root_state return root_state
# Deserialize root state from disk. # Deserialize root state from disk.
root_state = await self.load_state(root_state_token) root_state = await self.load_state(_substate_key(client_token, self.state))
# Create a new root state tree with all substates instantiated. # Create a new root state tree with all substates instantiated.
fresh_root_state = self.state(_reflex_internal_init=True) fresh_root_state = self.state(_reflex_internal_init=True)
if root_state is None: if root_state is None:
@ -2779,7 +2778,7 @@ class StateManagerDisk(StateManager):
else: else:
# Ensure all substates exist, even if they were not serialized previously. # Ensure all substates exist, even if they were not serialized previously.
root_state.substates = fresh_root_state.substates root_state.substates = fresh_root_state.substates
self.states[root_state_token] = root_state self.states[client_token] = root_state
await self.populate_substates(client_token, root_state, root_state) await self.populate_substates(client_token, root_state, root_state)
return root_state return root_state