From f42d1f4b0fb5dc2484444e221588a918f7eba838 Mon Sep 17 00:00:00 2001
From: Khaleel Al-Adhami <khaleel.aladhami@gmail.com>
Date: Thu, 16 Jan 2025 14:34:35 -0800
Subject: [PATCH] fix component state

---
 reflex/compiler/compiler.py | 13 ++++---------
 reflex/state.py             |  9 +++++++++
 reflex/vars/function.py     |  2 --
 3 files changed, 13 insertions(+), 11 deletions(-)

diff --git a/reflex/compiler/compiler.py b/reflex/compiler/compiler.py
index 894406431..c2f5662f4 100644
--- a/reflex/compiler/compiler.py
+++ b/reflex/compiler/compiler.py
@@ -560,18 +560,13 @@ def compile_unevaluated_page(
     """
     # Generate the component if it is a callable.
     component = page.component
-    component = (
-        component
-        if isinstance(component, Component)
-        else (Fragment.create(component) if isinstance(component, Var) else component())
-    )
-
-    # unpack components that return tuples in an rx.fragment.
-    if isinstance(component, tuple):
-        component = Fragment.create(*component)
 
     if isinstance(component, Var):
         component = Fragment.create(component)
+    elif isinstance(component, tuple):
+        component = Fragment.create(*component)
+    elif not isinstance(component, Component):
+        component = component()
 
     component._add_style_recursive(style or {}, theme)
 
diff --git a/reflex/state.py b/reflex/state.py
index 16875bec3..40b6dabf5 100644
--- a/reflex/state.py
+++ b/reflex/state.py
@@ -2554,6 +2554,9 @@ class ComponentState(State, mixin=True):
         Returns:
             A new instance of the Component with an independent copy of the State.
         """
+        from reflex.components import Component
+        from reflex.components.base.fragment import Fragment
+
         cls._per_component_state_instance_count += 1
         state_cls_name = f"{cls.__name__}_n{cls._per_component_state_instance_count}"
         component_state = type(
@@ -2565,6 +2568,12 @@ class ComponentState(State, mixin=True):
         # Save a reference to the dynamic state for pickle/unpickle.
         setattr(reflex.istate.dynamic, state_cls_name, component_state)
         component = component_state.get_component(*children, **props)
+        if isinstance(component, Var):
+            component = Fragment.create(component)
+        elif isinstance(component, tuple):
+            component = Fragment.create(*component)
+        elif not isinstance(component, Component):
+            component = component()
         component.State = component_state
         return component
 
diff --git a/reflex/vars/function.py b/reflex/vars/function.py
index f529df008..c3fe8f75f 100644
--- a/reflex/vars/function.py
+++ b/reflex/vars/function.py
@@ -722,7 +722,6 @@ class ArgsFunctionOperation(CachedVarOperation, FunctionVar[CALLABLE_TYPE]):
         cls,
         args_names: Sequence[Union[str, DestructuredArg]],
         return_expr: Var | Any,
-        /,
         default_values: Sequence[VarWithDefault | inspect.Parameter.empty] = (),
         rest: str | None = None,
         validators: Sequence[Callable[[Any], Optional[str]]] = (),
@@ -801,7 +800,6 @@ class ArgsFunctionOperationBuilder(
         cls,
         args_names: Sequence[Union[str, DestructuredArg]],
         return_expr: Var | Any,
-        /,
         default_values: Sequence[VarWithDefault | inspect.Parameter.empty] = (),
         rest: str | None = None,
         validators: Sequence[Callable[[Any], Optional[str]]] = (),