Fix type validation for browser storage annotations

news/+storage-annotation-hydration.bugfix.md

@@ -0,0 +1 @@
+Avoid false type errors when hydrating or assigning strings to vars annotated with `Cookie`, `LocalStorage`, or `SessionStorage`.

reflex/state.py

@@ -1303,7 +1303,13 @@ def _init_var(cls, name: str, prop: Var):
                 f'Found var "{prop._js_expr}" with type {prop._var_type}.'
             )
             raise VarTypeError(msg)
-        cls.__fields__[name]._var = prop
+        field = cls.__fields__[name]
+        field._var = prop
+        if types.safe_issubclass(field.type_, ClientStorageBase) or any(
+            issubclass(member, ClientStorageBase) for member in field._plain_types
+        ):
+            # Storage annotations configure persistence; browser values are strings.
+            field._plain_types |= {str}
         if cls.is_user_defined() and get_state_auto_setters() is True:
             cls._create_setter(name, prop)
         cls._set_default_value(name, prop)

tests/integration/tests_playwright/test_hydration_storage.py

@@ -1,5 +1,6 @@
 """Browser regressions for client-storage writes during hydration and assigned defaults."""
 
+import logging
 import re
 from collections.abc import Callable, Generator
 
@@ -20,8 +21,23 @@ class StorageState(rx.State):
         local: str = rx.LocalStorage("local-default", name="hydrate-local")
         session: str = rx.SessionStorage("session-default", name="hydrate-session")
         cookie: str = rx.Cookie("cookie-default", name="hydrate-cookie")
+        annotated_local: rx.Field[rx.LocalStorage] = rx.field(
+            rx.LocalStorage("light", name="hydrate-annotated-local", sync=True)
+        )
+        annotated_session: rx.SessionStorage | None = rx.SessionStorage(
+            "light", name="hydrate-annotated-session"
+        )
+        annotated_cookie: rx.Cookie = rx.Cookie(
+            "light", name="hydrate-annotated-cookie"
+        )
         nonce: rx.Field[str] = rx.field(default_factory=lambda: uuid.uuid4().hex)
 
+        @rx.event
+        def change_annotated(self):
+            """Assign plain strings to storage-annotated vars."""
+            for name in ("annotated_local", "annotated_session", "annotated_cookie"):
+                setattr(self, name, "changed")
+
         @rx.var
         def checked(self) -> str:
             """Clear rejected values in each storage type.
@@ -188,6 +204,19 @@ def index():
         """
         return rx.box(
             rx.text(StorageState.checked, id="checked"),
+            rx.text(
+                StorageState.annotated_local,
+                "|",
+                StorageState.annotated_session,
+                "|",
+                StorageState.annotated_cookie,
+                id="annotated",
+            ),
+            rx.button(
+                "Change annotated",
+                on_click=StorageState.change_annotated,
+                id="change-annotated",
+            ),
             rx.text(rx.cond(rx.State.is_hydrated, "true", "false"), id="hydrated"),
             rx.text(ReconcileState.token_hash, id="token-hash"),
         )
@@ -297,6 +326,52 @@ def test_hydration_reconciles_storage_in_get_delta(
     page.wait_for_function("localStorage.getItem('hydrate-token-hash') === 'fresh'")
 
 
+def test_storage_annotations_hydrate_and_persist(
+    hydration_storage_app: AppHarness,
+    page: Page,
+    caplog: pytest.LogCaptureFixture,
+    request: pytest.FixtureRequest,
+):
+    """Storage-annotated vars accept browser strings without logging type errors.
+
+    Args:
+        hydration_storage_app: The running app.
+        page: A fresh browser page.
+        caplog: Captured server log records.
+        request: The request used to remove the log handler after the test.
+    """
+    field_logger = logging.getLogger("reflex_base.vars.base")
+    field_logger.addHandler(caplog.handler)
+    request.addfinalizer(lambda: field_logger.removeHandler(caplog.handler))
+    assert hydration_storage_app.frontend_url is not None
+    page.goto(hydration_storage_app.frontend_url)
+    expect(page.locator("#hydrated")).to_have_text("true")
+    expect(page.locator("#annotated")).to_have_text("light|light|light")
+    page.evaluate("""() => {
+        localStorage.setItem('hydrate-annotated-local', 'saved');
+        sessionStorage.setItem('hydrate-annotated-session', 'saved');
+        document.cookie = 'hydrate-annotated-cookie=saved; path=/';
+    }""")
+
+    page.reload()
+    expect(page.locator("#hydrated")).to_have_text("true")
+    expect(page.locator("#annotated")).to_have_text("saved|saved|saved")
+    page.locator("#change-annotated").click()
+    page.wait_for_function("""() =>
+        localStorage.getItem('hydrate-annotated-local') === 'changed' &&
+        sessionStorage.getItem('hydrate-annotated-session') === 'changed' &&
+        document.cookie.split('; ').includes('hydrate-annotated-cookie=changed')
+    """)
+    page.reload()
+    expect(page.locator("#hydrated")).to_have_text("true")
+    expect(page.locator("#annotated")).to_have_text("changed|changed|changed")
+    assert not [
+        record.getMessage()
+        for record in caplog.records
+        if record.name == field_logger.name and record.levelno >= logging.ERROR
+    ]
+
+
 # The browser storage vars of the assigned page, by element id suffix.
 VARS = ("local", "session", "cookie", "factory", "pref")
 

tests/units/test_state.py

@@ -79,6 +79,7 @@
     OnLoadInternalState,
     State,
     StateUpdate,
+    _apply_client_storage_vars,
     _frontend_schema,
     _load_events_for_page,
     all_base_state_classes,
@@ -5908,6 +5909,79 @@ def test_frontend_schema_leaves_out_backend_vars():
     )
 
 
+@pytest.mark.parametrize("value", ["", "dark"])
+@pytest.mark.parametrize(
+    "var_name",
+    [
+        "cookie",
+        "local",
+        "session",
+        "field_cookie",
+        "field_local",
+        "field_session",
+        "optional",
+        "storage_union",
+        "plain",
+    ],
+)
+async def test_storage_annotations_accept_browser_strings(
+    var_name: str,
+    value: str,
+    caplog: pytest.LogCaptureFixture,
+    clean_registration_context: RegistrationContext,
+):
+    """Browser strings satisfy storage annotations during hydration and later writes.
+
+    Args:
+        var_name: The browser storage var to update.
+        value: The value sent by the browser.
+        caplog: Captured log records.
+        clean_registration_context: An isolated state registry.
+    """
+
+    class StorageAnnotationState(State):
+        cookie: rx.Cookie = rx.Cookie("light", name="cookie", path="/settings")
+        local: rx.LocalStorage = rx.LocalStorage("light", name="local", sync=True)
+        session: rx.SessionStorage = rx.SessionStorage("light", name="session")
+        field_cookie: rx.Field[rx.Cookie] = rx.field(rx.Cookie("light"))
+        field_local: rx.Field[rx.LocalStorage] = rx.field(rx.LocalStorage("light"))
+        field_session: rx.Field[rx.SessionStorage] = rx.field(
+            rx.SessionStorage("light")
+        )
+        plain: str = rx.LocalStorage("light", name="plain")
+        optional: rx.Field[rx.LocalStorage | None] = rx.field(rx.LocalStorage("light"))
+        storage_union: rx.Cookie | rx.SessionStorage = rx.Cookie("light")
+
+    root = State()
+    state = await root.get_state(StorageAnnotationState)
+    declared = StorageAnnotationState.get_fields()[var_name].default
+    assert isinstance(declared, (rx.Cookie, rx.LocalStorage, rx.SessionStorage))
+    options = declared.options()
+
+    with caplog.at_level(logging.ERROR):
+        applied = await _apply_client_storage_vars(
+            root,
+            {
+                f"{StorageAnnotationState.get_full_name()}.{var_name}{FIELD_MARKER}": value
+            },
+        )
+        assert applied == [(state, var_name)]
+        assert getattr(state, var_name) == value
+        assert var_name in state.dirty_vars
+        assert not caplog.records
+
+        setattr(state, var_name, "changed")
+        assert getattr(state, var_name) == "changed"
+        assert not caplog.records
+
+        setattr(state, var_name, 42)
+        assert len(caplog.records) == 1
+        assert "Expected field" in caplog.records[0].getMessage()
+
+    assert StorageAnnotationState.get_fields()[var_name].default is declared
+    assert declared.options() == options
+
+
 def test_reset_client_storage_uses_declared_factories():
     """Hydration evaluates declared factories for all browser storage types."""
     calls = []