Add mixin resolution to state lookup and token handling

news/+get-state-mixin.feature.md

@@ -0,0 +1 @@
+`get_state()` accepts a state mixin and returns the state it was mixed into, so shared handlers can reach it without naming the concrete class: `await self.get_state(AuthMixin)`. A `BaseStateToken` whose `cls` is a mixin addresses that state in the state manager too. Both raise `StateValueError` when no state, or more than one unrelated state, uses the mixin.

packages/reflex-base/news/+registry-mixin-lookup.feature.md

@@ -0,0 +1 @@
+Add `RegistrationContext.get_states_implementing()`, which lists the registered states that inherit from a state mixin, and `RegistrationContext.resolve_implementation()`, which returns the single state the mixin was mixed into and raises `StateValueError` when there is none or the choice is ambiguous.

packages/reflex-base/src/reflex_base/registry.py

@@ -3,7 +3,7 @@
 from __future__ import annotations
 
 import dataclasses
-from typing import TYPE_CHECKING, Any, Self
+from typing import TYPE_CHECKING, Any, Self, TypeVar
 
 from reflex_base.context.base import BaseContext
 from reflex_base.utils.exceptions import ReflexRuntimeError, StateValueError
@@ -19,6 +19,8 @@
     from reflex_base.utils.types import ArgsSpec
     from reflex_base.vars.base import Var
 
+STATE_TYPE = TypeVar("STATE_TYPE", bound="BaseState")
+
 
 def _default_bundled_libraries() -> list[str]:
     """Return the initial set of bundled libraries for a new context.
@@ -61,6 +63,19 @@ class RegistrationContext(BaseContext):
         default_factory=dict,
         repr=False,
     )
+    # Memoized lookups, invalidated whenever a base state is registered.
+    _implementations: dict[type[BaseState], tuple[type[BaseState], ...]] = (
+        dataclasses.field(
+            default_factory=dict,
+            repr=False,
+        )
+    )
+    _resolved_implementations: dict[type[BaseState], type[BaseState]] = (
+        dataclasses.field(
+            default_factory=dict,
+            repr=False,
+        )
+    )
     _config: Config | None = dataclasses.field(default=None, repr=False)
     decorated_pages: list[tuple[Callable, dict[str, Any]]] = dataclasses.field(
         default_factory=list,
@@ -215,6 +230,8 @@ def _register_base_state(self, state_cls: type[BaseState]) -> type[BaseState]:
             The registered base state class.
         """
         self.base_states[state_cls.get_full_name()] = state_cls
+        self._implementations.clear()
+        self._resolved_implementations.clear()
         for event_handler in state_cls.event_handlers.values():
             self._register_event_handler(event_handler, states=(state_cls,))
         if (parent_state := state_cls.get_parent_state()) is not None:
@@ -287,3 +304,66 @@ def get_substates(
         return self.base_state_substates.setdefault(
             base_state_cls.get_full_name(), set()
         )
+
+    def get_states_implementing(
+        self, interface: type[STATE_TYPE]
+    ) -> tuple[type[STATE_TYPE], ...]:
+        """Get every registered state class that is a subclass of the given interface.
+
+        Args:
+            interface: The state mixin (or state class) to look up implementations of.
+
+        Returns:
+            The registered states implementing the interface, in registration order.
+        """
+        if (cached := self._implementations.get(interface)) is not None:
+            return cached  # pyright: ignore [reportReturnType]
+        implementations = tuple(
+            state_cls
+            for state_cls in self.base_states.values()
+            if issubclass(state_cls, interface)
+        )
+        self._implementations[interface] = implementations
+        return implementations
+
+    def resolve_implementation(self, interface: type[STATE_TYPE]) -> type[STATE_TYPE]:
+        """Get the single registered state class the given interface was mixed into.
+
+        Args:
+            interface: The state mixin (or state class) to resolve.
+
+        Returns:
+            The state class the interface was mixed into.
+
+        Raises:
+            StateValueError: If no or more than one state implements the interface.
+        """
+        if (cached := self._resolved_implementations.get(interface)) is not None:
+            return cached  # pyright: ignore [reportReturnType]
+        implementations = self.get_states_implementing(interface)
+        # Keep only where the interface entered the tree: an implementation
+        # inheriting from another implementation is not a candidate.
+        roots = [
+            state_cls
+            for state_cls in implementations
+            if not any(
+                other is not state_cls and issubclass(state_cls, other)
+                for other in implementations
+            )
+        ]
+        if not roots:
+            msg = (
+                f"No registered state implements {interface.__name__}. "
+                "Define a state which inherits from it before resolving it."
+            )
+            raise StateValueError(msg)
+        if len(roots) > 1:
+            names = ", ".join(sorted(state.get_full_name() for state in roots))
+            msg = (
+                f"{interface.__name__} is implemented by more than one state: {names}. "
+                "Resolving it is ambiguous, pass the wanted state class explicitly."
+            )
+            raise StateValueError(msg)
+        resolved = roots[0]
+        self._resolved_implementations[interface] = resolved
+        return resolved

reflex/istate/manager/__init__.py

@@ -7,11 +7,12 @@
 from abc import ABC, abstractmethod
 from collections.abc import AsyncIterator
 from types import TracebackType
-from typing import TYPE_CHECKING, Generic, TypedDict, overload
+from typing import TYPE_CHECKING, Generic, TypedDict, cast, overload
 
 from reflex_base import constants
 from reflex_base.config import get_config
 from reflex_base.event import Event
+from reflex_base.registry import RegistrationContext
 from reflex_base.utils.exceptions import InvalidStateManagerModeError
 from typing_extensions import ReadOnly, Unpack, deprecated
 
@@ -98,7 +99,7 @@ def create(cls):
 
     @staticmethod
     def _coerce_token(token: StateToken[TOKEN_TYPE] | str) -> StateToken[TOKEN_TYPE]:
-        """Convert a legacy string token to a StateToken if needed.
+        """Normalize a token: convert legacy string tokens and resolve state mixins.
 
         Args:
             token: The token, either a StateToken or legacy string.
@@ -111,6 +112,12 @@ def _coerce_token(token: StateToken[TOKEN_TYPE] | str) -> StateToken[TOKEN_TYPE]
             from reflex.state import State
 
             return BaseStateToken.from_legacy_token(token, root_state=State)  # type: ignore[return-value]
+        # The store is never keyed on a mixin: resolve it to the state it was mixed into.
+        if getattr(token.cls, "_mixin", False):
+            resolved = RegistrationContext.get().resolve_implementation(
+                cast("type[BaseState]", token.cls)
+            )
+            return token.with_cls(cast("type[TOKEN_TYPE]", resolved))
         return token
 
     @overload

reflex/state.py

@@ -1793,11 +1793,15 @@ async def get_state(self, state_cls: type[T_STATE]) -> T_STATE:
         Allows for arbitrary access to sibling states from within an event handler.
 
         Args:
-            state_cls: The class of the state.
+            state_cls: The class of the state, or a state mixin to resolve.
 
         Returns:
             The instance of state_cls associated with this state's client_token.
         """
+        # A mixin has no place in the state tree of its own: resolve it.
+        if state_cls._mixin:
+            state_cls = RegistrationContext.get().resolve_implementation(state_cls)
+
         # Fast case - if this state instance is already cached, get_substate from root state.
         try:
             return self._get_state_from_cache(state_cls)

tests/units/reflex_base/test_registry.py

@@ -424,6 +424,130 @@ def test_bundled_libraries_isolated_between_contexts():
         assert "some-extra-lib" not in ctx_b.bundled_libraries
 
 
+def test_resolve_implementation(clean_registration_context: RegistrationContext):
+    """resolve_implementation finds the state a mixin was mixed into.
+
+    Args:
+        clean_registration_context: A fresh, empty registration context.
+    """
+    from reflex.state import BaseState
+
+    class ResolveMixin(BaseState, mixin=True):
+        pass
+
+    class ResolveImpl(ResolveMixin, BaseState):
+        pass
+
+    assert clean_registration_context.resolve_implementation(ResolveMixin) is (
+        ResolveImpl
+    )
+    assert clean_registration_context.get_states_implementing(ResolveMixin) == (
+        ResolveImpl,
+    )
+
+
+def test_resolve_implementation_ignores_subclasses(
+    clean_registration_context: RegistrationContext,
+):
+    """Subclasses of an implementation do not make the resolution ambiguous.
+
+    Args:
+        clean_registration_context: A fresh, empty registration context.
+    """
+    from reflex.state import BaseState
+
+    class SubclassMixin(BaseState, mixin=True):
+        pass
+
+    class SubclassImpl(SubclassMixin, BaseState):
+        pass
+
+    class SubclassChild(SubclassImpl):
+        pass
+
+    assert set(clean_registration_context.get_states_implementing(SubclassMixin)) == {
+        SubclassImpl,
+        SubclassChild,
+    }
+    assert clean_registration_context.resolve_implementation(SubclassMixin) is (
+        SubclassImpl
+    )
+
+
+def test_resolve_implementation_without_implementation(
+    clean_registration_context: RegistrationContext,
+):
+    """Resolving a mixin no state implements raises StateValueError.
+
+    Args:
+        clean_registration_context: A fresh, empty registration context.
+    """
+    from reflex.state import BaseState
+
+    class UnusedMixin(BaseState, mixin=True):
+        pass
+
+    assert clean_registration_context.get_states_implementing(UnusedMixin) == ()
+    with pytest.raises(StateValueError, match="No registered state implements"):
+        clean_registration_context.resolve_implementation(UnusedMixin)
+
+
+def test_resolve_implementation_ambiguous(
+    clean_registration_context: RegistrationContext,
+):
+    """Resolving a mixin used by two unrelated states raises StateValueError.
+
+    Args:
+        clean_registration_context: A fresh, empty registration context.
+    """
+    from reflex.state import BaseState
+
+    class AmbiguousMixin(BaseState, mixin=True):
+        pass
+
+    class AmbiguousImpl1(AmbiguousMixin, BaseState):
+        pass
+
+    class AmbiguousImpl2(AmbiguousMixin, BaseState):
+        pass
+
+    with pytest.raises(StateValueError, match="implemented by more than one state"):
+        clean_registration_context.resolve_implementation(AmbiguousMixin)
+
+
+def test_resolve_implementation_cache_is_invalidated(
+    clean_registration_context: RegistrationContext,
+):
+    """Registering a state invalidates previously cached lookups.
+
+    Args:
+        clean_registration_context: A fresh, empty registration context.
+    """
+    from reflex.state import BaseState
+
+    class CacheMixin(BaseState, mixin=True):
+        pass
+
+    class CacheImpl(CacheMixin, BaseState):
+        pass
+
+    assert clean_registration_context.resolve_implementation(CacheMixin) is CacheImpl
+    # Repeated lookups are served from the cache.
+    assert clean_registration_context.get_states_implementing(
+        CacheMixin
+    ) is clean_registration_context.get_states_implementing(CacheMixin)
+
+    class CacheImpl2(CacheMixin, BaseState):
+        pass
+
+    assert set(clean_registration_context.get_states_implementing(CacheMixin)) == {
+        CacheImpl,
+        CacheImpl2,
+    }
+    with pytest.raises(StateValueError, match="implemented by more than one state"):
+        clean_registration_context.resolve_implementation(CacheMixin)
+
+
 def test_reset_compile_caches_empties_per_compile_maps(
     clean_registration_context: RegistrationContext,
 ):