rx.cached_var and rx.background don't work in mixins

[TimChild]

Describe the bug

Trying to mixin an rx.background method causes error that "only background tasks should use async with self".

Mixing in rx.cached_var doesn't do any caching and instead works like a normal rx.var.

To Reproduce
Steps to reproduce the behavior:

import logging
import reflex as rx

from reflex_test.templates import template

logger = logging.getLogger(__name__)
logger.setLevel(logging.DEBUG)


class BasicMixin(rx.Base):
    a: int = 0
    b: int = 0
    c: int = 0

    @rx.var
    def var_a(self) -> int:
        logger.debug('var_a called')
        return self.a+10

    @rx.cached_var
    def cached_a(self) -> int:
        logger.debug('cached_a called')
        return self.a+20

    def increment_a(self):
        self.a += 1

    async def increment_b(self):
        self.b += 1

    @rx.background
    async def increment_c(self):
        async with self:
            self.c += 1


class StateWithBasicMixin(BasicMixin, rx.State):
    pass

@template(route='/mixin_cached_or_background_issue', title='Mixin Cached or Background Issue')
def index() -> rx.Component:
    return rx.container(
        rx.vstack(
            rx.heading('Mixin Cached or Background Issue', size='5'),
            rx.grid(
                rx.text('Attr a:'),
                rx.text(StateWithBasicMixin.a),
                rx.text('Var a:'),
                rx.text(StateWithBasicMixin.var_a),
                rx.text('Cached a:'),
                rx.text(StateWithBasicMixin.cached_a),
                rx.text('Attr b:'),
                rx.text(StateWithBasicMixin.b),
                rx.text('Attr c:'),
                rx.text(StateWithBasicMixin.c),
                columns='5',
                rows='2',
                flow='column'
            ),
            rx.hstack(
                rx.button('Increment a', on_click=StateWithBasicMixin.increment_a),
                rx.button('Increment b', on_click=StateWithBasicMixin.increment_b),
                rx.button('Increment c', on_click=StateWithBasicMixin.increment_c),
            ),
        ),
        padding="2em",
    )

Expected behavior
On click of increment b or increment c expect that the rx.cached_var is not run. This is true if the cached var is moved to the State, but not if it lives in the mixin.
On click of increment c expect that the value increments like a normal background task would. Instead an error is raised "TypeError: Only background task should use async with self to modify state.". Works as expected if moved to the State.

Specifics (please complete the following information):

• Python Version: 3.12
• Reflex Version: 0.4.9
• OS: WSL2
• Browser (Optional): Chrome

Additional context
If it's not easy to implement, it would be good to get some compilation errors :)

Related PR that allowed rx.var to work in mixins #2351

[benedikt-bartscher]

Thanks for reporting, I will investigate and fix this

[benedikt-bartscher]

@TimChild your issues regarding cached vars with mixins are fixed in #3254

[TimChild]

@benedikt-bartscher Awesome!
I thought there was a chance that it would be quick if you knew exactly what to look for, but it didn't stand out to me. Nice work!

[benedikt-bartscher]

@TimChild thanks, if you have any other issues regarding mixins, feel free to ping me.

[TimChild]

@benedikt-bartscher Well, while it is fresh in your mind... I do have a question/issue that is fairly related to this. Obviously, don't feel obliged to spend significant time on this, but I just wonder whether you may have a different way of looking at this problem.

I am trying to figure out how get around the fact that I can't use async code within an rx.cached_var. The idea of the rx.cached_var is pretty much ideal except for the fact that I can't use any async code in there because under the hood it is a property.

Below is an example that illustrates the idea:

class Data(BaseModel):
    attr_a: str
    attr_b: str

database: dict[int, Data] = {}

def load_data(data_id: int) -> Data:
    # Stand in for an async load from db function
    return database.get(data_id, Data(attr_a='', attr_b=''))

def process(input_a: str):
    # Stand in for an async processing function that stores result in db and returns only the id
    data = Data(attr_a=input_a.title(), attr_b=input_a.upper())
    next_id = len(database)
    database[next_id] = data
    return next_id


class HandlerState(rx.State):
    input_a: str

    data_id: int

    async def do_stuff_on_click(self):
        # Some process that stores result in redis/database
        data_id = process(self.input_a)
        # Only store the id here to avoid storing whole object in state and serialization issues
        self.data_id = data_id


class DisplayMixin(rx.Base):
    @rx.cached_var
    def data_a_cached(self) -> str:
        data = load_data(self.data_id)
        return data.attr_a

    @rx.cached_var
    def data_b_cached(self) -> str:
        data = load_data(self.data_id)
        return data.attr_b


class DisplayState(DisplayMixin, HandlerState):
    pass


class AlternativeDisplayMixin(rx.Base):
    data_attr_a: str = ""
    data_attr_b: str = ""

    async def update_display_info(self):
        data = load_data(self.data_id)
        self.data_attr_a = f'Attribute A: {data.attr_a}'
        self.data_attr_b = f'Attribute B: {data.attr_b}'

class AlternativeDisplayState(AlternativeDisplayMixin, HandlerState):
    pass


@template(route='/async_cached_var_issues', title='Async Cached Var Issues')
def index() -> rx.Component:
    return rx.container(
        rx.card(
            rx.heading('Handler stuff', size='5'),
            rx.input(value=HandlerState.input_a, label='Input A', on_change=HandlerState.set_input_a),
            rx.button('Do Stuff', on_click=HandlerState.do_stuff_on_click),
        ),
        rx.hstack(
            rx.card(
                rx.heading('Display data via cached_vars', size='5'),
                rx.markdown(f'{DisplayState.data_a_cached}\n\n{DisplayState.data_b_cached}'),
            ),
            rx.card(
                rx.heading('Alternative Display of data', size='5'),
                rx.markdown(f'{AlternativeDisplayState.data_attr_a}\n\n{AlternativeDisplayState.data_attr_b}'),
                rx.button('Update Display Info', on_click=AlternativeDisplayState.update_display_info),
            )
        ),
    )

In the example code here, I have left the load_data as a sync functions so that it does run to show the idea, but in reality, the load_data function would be async.

My thinking is that the HandlerState handles doing stuff, and can just store some id of a calculated result to avoid any serialization issues and/or storing a large amount of data in the reflex state.

Then displaying of the information is separated into a mixin because there are several ways I might want to display the processed data. The idea is that they can load the relevant data based on id and extract whatever should be shown on the frontend.

The DisplayMixin does this via cached_vars (which will at least work for sync code thanks to your fix!).
Positives: the display automatically updates if the data_id changes, and it doesn't need to store any more info in the State.
Negatives: Can't do async loading of data in the cached_var... Also, load needs to be called per cached_var (but that could itself be cached)

The AlternativeDisplayMixin does this via another event handler and setting new attrs on the State.
Positives: The data only needs to be loaded once and any number of things extracted from it
Negatives: The update needs to be triggered by another event, and it ends up storing all the data in the State (which I was trying to avoid).

I feel like there must be some nice clean solution to this... but it eludes me... The cached_var option feels right, except for the very obvious lack of async...