StateManager.get_state interface implementations differ in parameters.

[abulvenz]

Describe the bug
When using a middleware you can never be sure which implementation of StateManager you can expect.
So changes like removing redis_url as config parameter in rxconfig.py can break the system.

class MyMiddleware(rx.Middleware):
    async def postprocess(
        self,
        app: rx.App,
        state: rx.State,
        event: rx.state.Event,
        update: rx.state.StateUpdate,
    ) -> rx.state.StateUpdate | None:
        parent_state: AuthState = await app.state_manager.get_state(
            f"{event.token}_state.my_state", get_substates=False
        )

To Reproduce
Use the above code with and withouth redis_url

Expected behavior
No error. Maybe the args can be added to the StateManager.get_state function without functionality to have a consistent interface.

Specifics (please complete the following information):

• Python Version: 3.12
• Reflex Version: main
• OS: arch
• Browser (Optional): FF

Additional context
I stranded in this region, because due to state-sharding I do no longer get the entire state in the middleware as function parameters. So I need to retrieve the specific state from the state_manager as shown above. Since when using redis I do not need the substates which is more performant, I do not retrieve them.

@masenf Maybe you have an opinion on/solution for this, since you did thankfully implement the state-sharding feature.

[masenf]

i think the long term solution is the new get_state API of rx.BaseState: #2678

Essentially you'll be able to do auth_state = await state.get_state(AuthState) and it will work with both types of state manager.

It does not let you elect to not fetch the substates though, we need to get the substates so that computed vars in the substate which depend on the requested state will still update appropriately.

If fetching the substates is a performance-sensitive operation, then I would suggest flattening the state tree and having as few nested substates as possible.

[abulvenz]

Thanks for the outlook into the bright future.
I know that I have used an internal API that might be due to unexpected changes.
Maybe the StateManagerclasses could be streamlined in a way that both look more or less like a dict[token+substate_key, state] from the outside, i.e. that both are sharded. That could also reduce the internal overhead for retrieving states. I will look into your PR and use my current workaround. Closing this question. Thanks @masenf