Data-bound chart components 4/8: serve composite plan tokens

python/reflex_xy/app.py

@@ -67,6 +67,7 @@ def wire(namespace: XYNamespace) -> None:
     """Point the registry's fan-out seams at a namespace (setup and tests)."""
     registry.on_publish(namespace.broadcast_payload)
     registry.on_push(namespace.broadcast_message)
+    registry.on_error(namespace.broadcast_error)
 
 
 async def _xy_lifespan() -> None:

python/reflex_xy/namespace.py

@@ -39,14 +39,16 @@
 import asyncio
 import urllib.parse
 from collections.abc import Awaitable, Callable, Sequence
+from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, Optional
 
 from socketio import AsyncNamespace
 
 from xy.channel import handle_message
 
+from .plan import PlanError, PlanMissError
 from .registry import FigureEntry, FigureRegistry
-from .tokens import parse_token
+from .tokens import parse_plan_token, parse_token
 
 if TYPE_CHECKING:
     from xy._figure import Figure
@@ -104,9 +106,45 @@ def _handle_entry_message(entry: FigureEntry, content: Any) -> Any:
 
 # An async callable(token) -> Figure | None: given a parseable figure token,
 # rebuild the figure from Reflex state (wired by app.setup; see state_bridge).
+# May raise PlanError subclasses for spec-aware err frames.
 RebuildHook = Callable[[str], Awaitable[Optional["Figure"]]]
 
 
+@dataclass(frozen=True)
+class _RebuildFailure:
+    """Client-facing outcome of a failed rebuild attempt."""
+
+    error: str
+    resync: bool = False
+
+
+@dataclass(frozen=True)
+class _TokenIdentity:
+    """What a figure token reveals: session affinity and rebuildability."""
+
+    affinity_client: Optional[str]
+    rebuildable: bool
+    plan_bound: Optional[tuple[str, str]] = None  # (data_token, digest)
+
+
+def _token_identity(token: str) -> _TokenIdentity:
+    composite = parse_plan_token(token)
+    if composite is not None:
+        return _TokenIdentity(
+            affinity_client=composite.data.client_token,
+            rebuildable=True,
+            plan_bound=(composite.data_token, composite.digest),
+        )
+    parsed = parse_token(token)
+    if parsed is not None and parsed.kind == "figure":
+        return _TokenIdentity(affinity_client=parsed.client_token, rebuildable=True)
+    if parsed is not None:
+        # A bare data token names columns, never a figure: enforce affinity
+        # (it embeds a session) but never serve or rebuild it as one.
+        return _TokenIdentity(affinity_client=parsed.client_token, rebuildable=False)
+    return _TokenIdentity(affinity_client=None, rebuildable=False)
+
+
 def _plain(value: Any) -> Any:
     """Best-effort JSON-safe copy for small reply metadata.
 
@@ -159,7 +197,7 @@ def __init__(
         # room fan-out, so one cancelled waiter cannot cancel the attempt and
         # a failed builder is not rerun serially by every existing waiter.
         self._rebuild_attempts: dict[
-            str, asyncio.Task[tuple[Optional[FigureEntry], Optional[str]]]
+            str, asyncio.Task[tuple[Optional[FigureEntry], Optional[_RebuildFailure]]]
         ] = {}
 
     # -- connection lifecycle ------------------------------------------------
@@ -208,7 +246,13 @@ async def on_sub(self, sid: str, data: Any) -> None:
                 # path has proved the SID is still live. A concurrent
                 # disconnect can now only remove this record, never precede
                 # and be undone by it.
-                self.registry.subscribe(token, sid, rebuildable=parse_token(token) is not None)
+                identity = _token_identity(token)
+                self.registry.subscribe(token, sid, rebuildable=identity.rebuildable)
+                if identity.plan_bound is not None:
+                    # Index the mounted plan before serving, so a column
+                    # republish racing this subscribe rebuilds it (the
+                    # re-read below then serves that fresher generation).
+                    self.registry.bind_plan(*identity.plan_bound)
                 # A normal state publish can replace a just-rebuilt entry
                 # while its room-wide broadcast is still completing, before
                 # this SID joins. Re-read after the join: replacements before
@@ -438,7 +482,7 @@ def _release_subscription_lock(self, key: tuple[str, str], lock: asyncio.Lock) -
 
     def _start_rebuild_attempt(
         self, token: str
-    ) -> asyncio.Task[tuple[Optional[FigureEntry], Optional[str]]]:
+    ) -> asyncio.Task[tuple[Optional[FigureEntry], Optional[_RebuildFailure]]]:
         """Start and retain one shared cache-miss attempt for ``token``."""
         # Install the guard before publishing the task in ``_rebuild_attempts``.
         # A concurrent waiter can then distinguish this live attempt from one
@@ -447,7 +491,9 @@ def _start_rebuild_attempt(
         task = asyncio.create_task(self._run_rebuild_attempt(token, entry, guard))
         self._rebuild_attempts[token] = task
 
-        def forget(completed: asyncio.Task[tuple[Optional[FigureEntry], Optional[str]]]) -> None:
+        def forget(
+            completed: asyncio.Task[tuple[Optional[FigureEntry], Optional[_RebuildFailure]]],
+        ) -> None:
             if self._rebuild_attempts.get(token) is completed:
                 self._rebuild_attempts.pop(token, None)
 
@@ -459,19 +505,26 @@ async def _run_rebuild_attempt(
         token: str,
         entry: Optional[FigureEntry],
         guard: Optional[object],
-    ) -> tuple[Optional[FigureEntry], Optional[str]]:
+    ) -> tuple[Optional[FigureEntry], Optional[_RebuildFailure]]:
         """Build, conditionally publish, and fan out one total rebuild attempt."""
+        unknown = _RebuildFailure("unknown figure token")
         if entry is not None:
             return entry, None
         if guard is None:  # defensive: a miss always receives one bounded guard
-            return None, "unknown figure token"
+            return None, unknown
 
         try:
             rebuild = self._rebuild
             if rebuild is None:
-                return None, "unknown figure token"
+                return None, unknown
             try:
                 figure = await rebuild(token)
+            except PlanError as exc:
+                # Spec-aware failures of the data-bound tier: a stale plan
+                # digest asks the client to resync (the recompiled page
+                # carries the new digest); a bind mismatch names both sides
+                # and is not retryable as-is.
+                return None, _RebuildFailure(str(exc), resync=isinstance(exc, PlanMissError))
             except Exception:  # noqa: BLE001 - user builder code is an input boundary
                 figure = None
 
@@ -480,12 +533,12 @@ async def _run_rebuild_attempt(
                 # awaiting. Use it instead of reporting a stale rebuild failure.
                 entry = self.registry.get(token)
                 if entry is None:
-                    return None, "unknown figure token"
+                    return None, unknown
                 return entry, None
 
             entry, inserted = self.registry.publish_if_missing(token, figure, guard=guard)
             if entry is None:
-                return None, "unknown figure token"
+                return None, unknown
             if not inserted:
                 return entry, None
 
@@ -499,15 +552,15 @@ async def _run_rebuild_attempt(
                     current = self.registry.get(token)
                     if current is not None:
                         return current, None
-                return None, "rebuild failed"
+                return None, _RebuildFailure("rebuild failed")
 
             # ``broadcast_payload`` intentionally no-ops when its generation
             # went stale. Resolve that race explicitly: a replacement wins,
             # while a concurrent release becomes a closed error and retry.
             if not self.registry.is_current(token, entry):
                 current = self.registry.get(token)
                 if current is None:
-                    return None, "unknown figure token"
+                    return None, unknown
                 return current, None
             return entry, None
         finally:
@@ -527,10 +580,10 @@ async def _entry_for(
         token = self._token_of(data)
         if token is None:
             return None, None, False
-        parsed = parse_token(token)
-        if parsed is not None:
+        identity = _token_identity(token)
+        if identity.affinity_client is not None:
             session = await self.get_session(sid)
-            if session.get("client_token") != parsed.client_token:
+            if session.get("client_token") != identity.affinity_client:
                 await self._err(sid, token, "figure belongs to another session")
                 return token, None, False
         entry, rebuild_guarded = self.registry.get_with_rebuild_guard(token)
@@ -541,7 +594,12 @@ async def _entry_for(
         # forever for older user rebuild code. An entry whose guard is still
         # valid remains provisional until its rebuild fan-out completes.
         initially_missing = entry is None or (attempt is not None and rebuild_guarded)
-        if parsed is not None and allow_rebuild and self._rebuild is not None and initially_missing:
+        if (
+            identity.rebuildable
+            and allow_rebuild
+            and self._rebuild is not None
+            and initially_missing
+        ):
             # The task spans builder, conditional insertion, and existing-room
             # fan-out. All requests that observed this in-flight miss share its
             # result and must drop pre-payload interactions, even when another
@@ -553,14 +611,27 @@ async def _entry_for(
                 attempt = None
             if attempt is None:
                 attempt = self._start_rebuild_attempt(token)
-            entry, error = await asyncio.shield(attempt)
-            if error is not None:
-                await self._err(sid, token, error)
+            entry, failure = await asyncio.shield(attempt)
+            if failure is not None:
+                await self._err(sid, token, failure.error, resync=failure.resync)
                 return token, None, initially_missing
         if entry is None:
             await self._err(sid, token, "unknown figure token")
             return token, None, False
         return token, entry, initially_missing
 
-    async def _err(self, sid: str, token: Optional[str], error: str) -> None:
-        await self.emit("err", {"fig": token, "error": error}, to=sid)
+    async def _err(
+        self, sid: str, token: Optional[str], error: str, *, resync: bool = False
+    ) -> None:
+        envelope: dict[str, Any] = {"fig": token, "error": error}
+        if resync:
+            envelope["resync"] = True
+        await self.emit("err", envelope, to=sid)
+
+    async def broadcast_error(self, token: str, error: str, resync: bool = False) -> None:
+        """Room-wide err frame for server-side failures with no request to
+        answer (e.g. a column republish whose plan bind fails)."""
+        envelope: dict[str, Any] = {"fig": token, "error": error}
+        if resync:
+            envelope["resync"] = True
+        await self.emit("err", envelope, room=self._room(token))

python/reflex_xy/registry.py

@@ -125,8 +125,10 @@ def __init__(self, ttl_seconds: float = DEFAULT_TTL_SECONDS) -> None:
         # and the data-token -> {plan digest} index that lets a column
         # republish rebuild + broadcast every mounted dependent figure. The
         # index is bounded by mounted plans: entries are added when a
-        # composite figure binds and dropped when a republish finds neither
-        # a cached figure nor live subscribers under the composite token.
+        # composite figure binds (namespace sub) and dropped by
+        # _unbind_plan_if_unmounted_locked on every transition that can end
+        # the mount — unsubscribe, disconnect, release, failed-rebuild
+        # cleanup, TTL sweep, and a republish that finds it unmounted.
         self._column_entries: dict[str, ColumnEntry] = {}
         self._digests_by_data_token: dict[str, set[str]] = {}
         # A mounted client remains in its socket room when a TTL sweep evicts
@@ -458,6 +460,7 @@ def release(self, token: str) -> None:
             self._invalidate_rebuild_guards_locked(token)
             entry = self._entries.pop(token, None)
             self._retain_removed_version_locked(token, entry)
+            self._unbind_plan_if_unmounted_locked(token)
 
     def remove_if_current(self, token: str, expected: FigureEntry, *, guard: object) -> bool:
         """Remove ``expected`` only while its rebuild guard is still valid.
@@ -474,11 +477,35 @@ def remove_if_current(self, token: str, expected: FigureEntry, *, guard: object)
                 return False
             del self._entries[token]
             self._retain_removed_version_locked(token, expected)
+            self._unbind_plan_if_unmounted_locked(token)
             return True
 
     def _invalidate_rebuild_guards_locked(self, token: str) -> None:
         self._active_rebuild_guards.pop(token, None)
 
+    def _unbind_plan_if_unmounted_locked(self, token: str) -> None:
+        """Drop a composite plan token's binding once nothing serves or
+        watches it (mutex held; no-op for non-plan tokens).
+
+        This is the other half of the index invariant "bounded by mounted
+        plans": ``bind_plan`` inserts on subscribe, and every transition that
+        can end a mount — unsubscribe, disconnect, release, failed-rebuild
+        cleanup, the TTL sweep — funnels here, so short-lived sessions cannot
+        accumulate bindings that only a republish would have collected.
+        """
+        from .tokens import parse_plan_token
+
+        parsed = parse_plan_token(token)
+        if parsed is None:
+            return
+        if token in self._entries or self._rebuildable_subscribers.get(token):
+            return
+        digests = self._digests_by_data_token.get(parsed.data_token)
+        if digests is not None:
+            digests.discard(parsed.digest)
+            if not digests:
+                self._digests_by_data_token.pop(parsed.data_token, None)
+
     def _retain_removed_version_locked(self, token: str, entry: Optional[FigureEntry]) -> None:
         if self._rebuildable_subscribers.get(token):
             version = 0 if entry is None else entry.version
@@ -518,6 +545,7 @@ def _unsubscribe_locked(self, token: str, sid: str) -> None:
             if not subscribers:
                 self._rebuildable_subscribers.pop(token, None)
                 self._evicted_versions.pop(token, None)
+                self._unbind_plan_if_unmounted_locked(token)
 
         tokens = self._rebuildable_tokens_by_sid.get(sid)
         if tokens is not None:
@@ -537,6 +565,7 @@ def disconnect(self, sid: str) -> None:
                 if not subscribers:
                     self._rebuildable_subscribers.pop(token, None)
                     self._evicted_versions.pop(token, None)
+                    self._unbind_plan_if_unmounted_locked(token)
 
     def tokens(self) -> list[str]:
         with self._mutex:
@@ -613,13 +642,9 @@ def _rebuild_dependent(self, data_token: str, digest: str, column_entry: ColumnE
             )
             if not mounted:
                 # Nothing serves or watches this plan anymore: forget the
-                # binding (this is what keeps the index bounded by mounts);
-                # a later subscribe rebuilds from scratch and re-indexes.
-                digests = self._digests_by_data_token.get(data_token)
-                if digests is not None:
-                    digests.discard(digest)
-                    if not digests:
-                        self._digests_by_data_token.pop(data_token, None)
+                # binding (the same invariant every unmount transition
+                # enforces); a later subscribe rebuilds and re-indexes.
+                self._unbind_plan_if_unmounted_locked(composite)
                 return
         parsed = parse_token(data_token)
         source = (
@@ -897,6 +922,7 @@ def sweep(self, *, now: Optional[float] = None) -> list[str]:
                     self._invalidate_rebuild_guards_locked(token)
                     self._retain_removed_version_locked(token, entry)
                     del self._entries[token]
+                    self._unbind_plan_if_unmounted_locked(token)
                     dropped.append(token)
             for token, column_entry in list(self._column_entries.items()):
                 if now - column_entry.last_access > self._ttl: