Serve replicator-mode flag checks only once the cache is ready

README.md

@@ -931,6 +931,16 @@ When running in Replicator Mode, the client will:
 - Use cached data populated by the external replicator service
 - Fall back to direct API calls if the replicator is not available
 
+#### Cache Readiness
+
+The SDK serves flag checks from the shared cache only once the replicator reports that the cache is ready. The health endpoint answers `503` with `"ready": false` while the replicator is still loading the cache for its `cache_version`, and `200` with `"ready": true` once that load is complete. Until the SDK has seen `"ready": true`, `check_flag`, `check_flag_with_entitlement` and `check_flags` skip the cache and call the Schematic API, falling back to the flag default if the API call fails. Once the cache is ready, single and bulk checks both evaluate locally from the cache, and a flag that is not in the cache still falls back to the API.
+
+Readiness changes only when the replicator answers with a `"ready"` field. If a health check fails (connection error, timeout, or an unreadable response), the SDK keeps the readiness and `cache_version` from the last answer it read, so a replicator restart or a missed poll does not send every check to the API. Before the first answer, the cache is not ready.
+
+The same gate applies to `check()` with `usage` in client credit-lease mode: while the cache is not ready, it runs as a plain check through the API instead of gating on a local lease. `prewarm` is not a flag check and still looks up companies through the DataStream without waiting for readiness.
+
+`DataStreamClient.is_cache_ready()` reports this readiness. `is_connected()` returns the same value in Replicator Mode and is kept for backward compatibility.
+
 ## Contributing
 
 While we value open-source contributions to this SDK, this library is generated programmatically.

src/schematic/client.py

@@ -1365,6 +1365,21 @@ async def initialize(self) -> None:
     def _get_datastream(self) -> Optional[DataStreamClient]:
         return self._datastream_client
 
+    def _get_flag_check_datastream(self) -> Optional[DataStreamClient]:
+        """The DataStream client flag checks may evaluate against, or None to
+        use the API.
+
+        Single and bulk flag checks both go through this gate so they cannot
+        drift. In replicator mode the shared cache is read only once the
+        replicator reports it ready (``is_cache_ready()``); before that the
+        check takes the API path, which falls back to the flag default if the
+        API fails.
+        """
+        ds = self._get_datastream()
+        if ds is None or not ds.is_cache_ready():
+            return None
+        return ds
+
     async def check_flag(
         self,
         flag_key: str,
@@ -1386,7 +1401,7 @@ async def check_flag_with_entitlement(
             return self._default_response(flag_key, options, REASON_OFFLINE)
 
         # Try DataStream first if available
-        ds = self._get_datastream()
+        ds = self._get_flag_check_datastream()
         if ds is not None:
             try:
                 resp = await ds.check_flag(
@@ -1414,8 +1429,10 @@ async def check_flags(
 
         # DataStream evaluation only makes sense when specific keys are
         # requested AND the client is connected — the "give me everything"
-        # semantic only exists via the bulk API.
-        ds = self._get_datastream()
+        # semantic only exists via the bulk API. The readiness gate is shared
+        # with check_flag; in replicator mode is_connected() is the same cache
+        # readiness, so only websocket mode additionally needs a live socket.
+        ds = self._get_flag_check_datastream()
         if ds is not None and ds.is_connected() and flag_keys:
             try:
                 results: List[CheckFlagResponseData] = []
@@ -1672,7 +1689,9 @@ async def enqueue_flag_check(body: EventBodyFlagCheck) -> None:
 
         return await check_with_lease(
             CreditCheckDeps(
-                datastream=self._datastream_client,
+                # Same readiness gate as the plain checks: with no DataStream
+                # the lease flow defers to the plain check, which uses the API.
+                datastream=self._get_flag_check_datastream(),
                 lease_store=lease_store,
                 reservations=reservations,
                 manager=manager,
@@ -1954,10 +1973,11 @@ async def _update_company_metrics(
         self, company: Optional[Dict[str, str]], event: str, quantity: Optional[int],
     ) -> None:
         # Not gated on is_connected(). Flag checks keep evaluating from the
-        # cache while a replicator reports not ready or the WebSocket is down,
-        # so the cached metric has to keep counting the usage tracked in the
-        # meantime or those checks would enforce limits against a frozen
-        # figure. The bump cannot double count: the server's next push for the
+        # cache while the WebSocket is down, and resume from the shared cache
+        # once a replicator reports ready again, so the cached metric has to
+        # keep counting the usage tracked in the meantime or those checks
+        # would enforce limits against a frozen figure. The bump cannot double
+        # count: the server's next push for the
         # company replaces the metric outright, and a company missing from the
         # cache is left alone.
         ds = self._get_datastream()

src/schematic/datastream/datastream_client.py

@@ -314,10 +314,33 @@ async def start(self) -> None:
         self._ws_client.start()
 
     def is_connected(self) -> bool:
+        """Whether the DataStream can serve data.
+
+        In websocket mode this reports whether the websocket is connected. In
+        replicator mode it reports replicator readiness, the same value as
+        ``is_cache_ready()``, which is the clearer name for that question.
+        """
         if self._replicator_mode:
             return self._replicator_ready
         return self._ws_client.is_connected() if self._ws_client else False
 
+    def is_cache_ready(self) -> bool:
+        """Whether flag checks may be evaluated from the cache.
+
+        In replicator mode this is the ``ready`` field of the last health
+        response the SDK could read: true once the replicator has finished
+        loading the cache for its current ``cache_version``, false before
+        that. A failed poll leaves it unchanged. While it is false, ``AsyncSchematic`` does not
+        read the cache and answers flag checks from the API.
+
+        Outside replicator mode the SDK fills its own cache over the WebSocket
+        and fetches what it lacks on demand, so there is nothing to wait for
+        and this returns true.
+        """
+        if self._replicator_mode:
+            return self._replicator_ready
+        return True
+
     def is_replicator_ready(self) -> bool:
         return self._replicator_ready
 
@@ -1177,42 +1200,53 @@ async def _replicator_health_loop(self) -> None:
                 break
 
     async def _check_replicator_health(self) -> None:
+        """Poll the replicator's health URL and record readiness and cache version.
+
+        Readiness changes only on an explicit answer from the replicator: a
+        parsed body whose ``ready`` field is a boolean, whatever the HTTP
+        status. Before the replicator has loaded the cache it answers 503 with
+        ``ready: false`` and the ``cache_version`` it is loading.
+
+        A failed poll (connection error, timeout, unparseable body, or a body
+        without a boolean ``ready``) changes nothing: it says nothing about
+        whether the shared cache is complete, and flipping every check to the
+        API on one missed poll would turn a replicator blip into a burst of
+        API checks across every SDK instance.
+        """
         if not self._replicator_health_url:
             return
         try:
             if not self._health_check_client:
                 self._health_check_client = httpx.AsyncClient(timeout=REPLICATOR_HEALTH_TIMEOUT_S)
             resp = await self._health_check_client.get(self._replicator_health_url)
-            resp.raise_for_status()
             health_data = resp.json()
+            if not isinstance(health_data, dict) or not isinstance(health_data.get("ready"), bool):
+                raise ValueError(f"unexpected health response body: {health_data!r}")
+        except Exception as exc:
+            self._logger.debug("Replicator health check failed, keeping last readiness: %s", exc)
+            return
 
-            self._logger.debug("Replicator health response: %s", health_data)
-
-            was_ready = self._replicator_ready
-            self._replicator_ready = health_data.get("ready", False)
+        self._logger.debug("Replicator health response (%s): %s", resp.status_code, health_data)
 
-            new_version = health_data.get("cache_version") or health_data.get("cacheVersion")
-            if new_version and new_version != self._replicator_cache_version:
-                old = self._replicator_cache_version
-                self._replicator_cache_version = new_version
-                self._logger.info("Cache version changed from %s to %s", old, new_version)
+        new_version = health_data.get("cache_version") or health_data.get("cacheVersion")
+        if new_version and new_version != self._replicator_cache_version:
+            old = self._replicator_cache_version
+            self._replicator_cache_version = new_version
+            self._logger.info("Cache version changed from %s to %s", old, new_version)
 
-            if self._replicator_ready and not was_ready:
-                self._logger.info("External replicator is now ready")
-                if self._on_replicator_health_changed:
-                    self._on_replicator_health_changed(True)
-            elif not self._replicator_ready and was_ready:
-                self._logger.info("External replicator is no longer ready")
-                if self._on_replicator_health_changed:
-                    self._on_replicator_health_changed(False)
+        self._set_replicator_ready(health_data["ready"])
 
-        except Exception as exc:
-            if self._replicator_ready:
-                self._replicator_ready = False
-                self._logger.info("External replicator is no longer ready")
-                if self._on_replicator_health_changed:
-                    self._on_replicator_health_changed(False)
-            self._logger.debug("Replicator health check failed: %s", exc)
+    def _set_replicator_ready(self, ready: bool) -> None:
+        was_ready = self._replicator_ready
+        self._replicator_ready = ready
+        if ready and not was_ready:
+            self._logger.info("External replicator is now ready")
+            if self._on_replicator_health_changed:
+                self._on_replicator_health_changed(True)
+        elif not ready and was_ready:
+            self._logger.info("External replicator is no longer ready")
+            if self._on_replicator_health_changed:
+                self._on_replicator_health_changed(False)
 
     async def get_replicator_cache_version_async(self, timeout_s: float = REPLICATOR_CACHE_VERSION_TIMEOUT_S) -> Optional[str]:
         """Attempt to fetch cache version immediately if not already available."""
@@ -1224,8 +1258,10 @@ async def get_replicator_cache_version_async(self, timeout_s: float = REPLICATOR
                 if not self._health_check_client:
                     self._health_check_client = httpx.AsyncClient(timeout=timeout_s)
                 resp = await self._health_check_client.get(self._replicator_health_url)
-                if resp.status_code == 200:
-                    data = resp.json()
+                # A 503 from a replicator still loading the cache carries the
+                # cache_version it is loading, so read the body on any status.
+                data = resp.json()
+                if isinstance(data, dict):
                     version = data.get("cache_version") or data.get("cacheVersion")
                     if version:
                         self._replicator_cache_version = version

tests/custom/test_client.py

@@ -2589,6 +2589,7 @@ def _lease_datastream(results: list, **overrides) -> ScriptedDataStream:
         **overrides,
     )
     datastream.is_connected = MagicMock(return_value=True)  # type: ignore[attr-defined]
+    datastream.is_cache_ready = MagicMock(return_value=True)  # type: ignore[attr-defined]
     datastream.close = AsyncMock()  # type: ignore[attr-defined]
     datastream.update_company_metrics = AsyncMock()  # type: ignore[attr-defined]
     return datastream
@@ -2657,6 +2658,23 @@ async def test_auto_with_datastream_gates_on_a_local_lease(self):
         finally:
             await self._drain(client)
 
+    async def test_replicator_cache_not_ready_checks_plainly_over_the_api(self):
+        """The lease check reads the cache, so it shares the flag checks'
+        readiness gate: while the replicator cache is not ready it runs as a
+        plain check against the API and takes no lease."""
+        client = _async_lease_client()
+        datastream = _lease_datastream([LEASE_PROBE, LEASE_GATE])
+        datastream.is_cache_ready = MagicMock(return_value=False)  # type: ignore[attr-defined]
+        client._datastream_client = datastream
+        try:
+            result = await self._check(client)
+            assert result.allowed is True
+            assert result.reservation is None
+            client.features.check_flag.assert_awaited_once()
+            client.credits.acquire_credit_lease.assert_not_awaited()
+        finally:
+            await self._drain(client)
+
     async def test_auto_falls_back_to_server_mode_when_datastream_fails_to_start(self):
         client = _async_lease_client()
         client._datastream_client.start = AsyncMock(side_effect=RuntimeError("no socket"))  # type: ignore[union-attr]