From 6cbdd0016985f0f7eadce725faf3bb2dddadbe49 Mon Sep 17 00:00:00 2001
From: Khaleel Al-Adhami <khaleel.aladhami@gmail.com>
Date: Tue, 11 Feb 2025 16:46:06 -0800
Subject: [PATCH] fix toast provider needed (#4801)

* fix toast provider needed

* fix tests
---
 reflex/app.py                      | 25 ++++++++++++++++++++++++-
 reflex/components/component.py     |  3 ---
 reflex/components/sonner/toast.py  |  6 +++++-
 reflex/components/sonner/toast.pyi |  8 +++++++-
 tests/units/test_app.py            |  6 ++++++
 5 files changed, 42 insertions(+), 6 deletions(-)

diff --git a/reflex/app.py b/reflex/app.py
index 281727b3f..d290b8f49 100644
--- a/reflex/app.py
+++ b/reflex/app.py
@@ -355,6 +355,9 @@ class App(MiddlewareMixin, LifespanMixin):
         [Exception], Union[EventSpec, List[EventSpec], None]
     ] = default_backend_exception_handler
 
+    # Put the toast provider in the app wrap.
+    bundle_toaster: bool = True
+
     @property
     def api(self) -> FastAPI | None:
         """Get the backend api.
@@ -1010,6 +1013,10 @@ class App(MiddlewareMixin, LifespanMixin):
         should_compile = self._should_compile()
 
         if not should_compile:
+            if self.bundle_toaster:
+                from reflex.components.sonner.toast import Toaster
+
+                Toaster.is_used = True
             with console.timing("Evaluate Pages (Backend)"):
                 for route in self._unevaluated_pages:
                     console.debug(f"Evaluating page: {route}")
@@ -1039,6 +1046,20 @@ class App(MiddlewareMixin, LifespanMixin):
             + adhoc_steps_without_executor,
         )
 
+        if self.bundle_toaster:
+            from reflex.components.component import memo
+            from reflex.components.sonner.toast import toast
+
+            internal_toast_provider = toast.provider()
+
+            @memo
+            def memoized_toast_provider():
+                return internal_toast_provider
+
+            toast_provider = Fragment.create(memoized_toast_provider())
+
+            app_wrappers[(1, "ToasterProvider")] = toast_provider
+
         with console.timing("Evaluate Pages (Frontend)"):
             for route in self._unevaluated_pages:
                 console.debug(f"Evaluating page: {route}")
@@ -1081,7 +1102,9 @@ class App(MiddlewareMixin, LifespanMixin):
             component = app_wrap(self._state is not None)
             if component is not None:
                 app_wrappers[key] = component
-                custom_components |= component._get_all_custom_components()
+
+        for component in app_wrappers.values():
+            custom_components |= component._get_all_custom_components()
 
         if self.error_boundary:
             console.deprecate(
diff --git a/reflex/components/component.py b/reflex/components/component.py
index 9466933c5..d27bddf78 100644
--- a/reflex/components/component.py
+++ b/reflex/components/component.py
@@ -1792,9 +1792,6 @@ class CustomComponent(Component):
             include_children=include_children, ignore_ids=ignore_ids
         )
         yield from filter(lambda prop: isinstance(prop, Var), self.props.values())
-        yield from self.get_component(self)._get_vars(
-            include_children=include_children, ignore_ids=ignore_ids
-        )
 
     @lru_cache(maxsize=None)  # noqa: B019
     def get_component(self) -> Component:
diff --git a/reflex/components/sonner/toast.py b/reflex/components/sonner/toast.py
index e215f356f..d1f9464d8 100644
--- a/reflex/components/sonner/toast.py
+++ b/reflex/components/sonner/toast.py
@@ -8,6 +8,7 @@ from reflex.base import Base
 from reflex.components.component import Component, ComponentNamespace
 from reflex.components.lucide.icon import Icon
 from reflex.components.props import NoExtrasAllowedProps, PropsBase
+from reflex.constants.base import Dirs
 from reflex.event import EventSpec, run_script
 from reflex.style import Style, resolved_color_mode
 from reflex.utils import format
@@ -27,7 +28,10 @@ LiteralPosition = Literal[
     "bottom-right",
 ]
 
-toast_ref = Var(_js_expr="refs['__toast']")
+toast_ref = Var(
+    _js_expr="refs['__toast']",
+    _var_data=VarData(imports={f"$/{Dirs.STATE_PATH}": [ImportVar(tag="refs")]}),
+)
 
 
 class ToastAction(Base):
diff --git a/reflex/components/sonner/toast.pyi b/reflex/components/sonner/toast.pyi
index 7ff0b9196..cb637bfff 100644
--- a/reflex/components/sonner/toast.pyi
+++ b/reflex/components/sonner/toast.pyi
@@ -9,9 +9,12 @@ from reflex.base import Base
 from reflex.components.component import Component, ComponentNamespace
 from reflex.components.lucide.icon import Icon
 from reflex.components.props import NoExtrasAllowedProps, PropsBase
+from reflex.constants.base import Dirs
 from reflex.event import EventSpec, EventType
 from reflex.style import Style
+from reflex.utils.imports import ImportVar
 from reflex.utils.serializers import serializer
+from reflex.vars import VarData
 from reflex.vars.base import Var
 
 LiteralPosition = Literal[
@@ -22,7 +25,10 @@ LiteralPosition = Literal[
     "bottom-center",
     "bottom-right",
 ]
-toast_ref = Var(_js_expr="refs['__toast']")
+toast_ref = Var(
+    _js_expr="refs['__toast']",
+    _var_data=VarData(imports={f"$/{Dirs.STATE_PATH}": [ImportVar(tag="refs")]}),
+)
 
 class ToastAction(Base):
     label: str
diff --git a/tests/units/test_app.py b/tests/units/test_app.py
index 5d25b09ac..88cb36509 100644
--- a/tests/units/test_app.py
+++ b/tests/units/test_app.py
@@ -1307,8 +1307,11 @@ def test_app_wrap_compile_theme(
         + "<RadixThemesColorModeProvider>"
         "<RadixThemesTheme accentColor={\"plum\"} css={{...theme.styles.global[':root'], ...theme.styles.global.body}}>"
         "<Fragment>"
+        "<MemoizedToastProvider/>"
+        "<Fragment>"
         "{children}"
         "</Fragment>"
+        "</Fragment>"
         "</RadixThemesTheme>"
         "</RadixThemesColorModeProvider>"
         + ("</StrictMode>" if react_strict_mode else "")
@@ -1371,8 +1374,11 @@ def test_app_wrap_priority(
         "<RadixThemesColorModeProvider>"
         "<Fragment2>"
         "<Fragment>"
+        "<MemoizedToastProvider/>"
+        "<Fragment>"
         "{children}"
         "</Fragment>"
+        "</Fragment>"
         "</Fragment2>"
         "</RadixThemesColorModeProvider>"
         "</RadixThemesText>"