Skip to content

Commit 14a787e

Browse files
committed
fix(telemetry): avoid double-counting inherited storage fields
State classes that inherit Cookie/LocalStorage/SessionStorage fields were re-counting the parent's fields on each descendant. Walk only the fields declared on each class so the per-feature counters reflect distinct storage vars across the state tree.
1 parent a6b877d commit 14a787e

4 files changed

Lines changed: 352 additions & 2 deletions

File tree

reflex/utils/telemetry_accounting.py

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -267,8 +267,9 @@ def _walk_state_features(
267267
for handler in state_cls.event_handlers.values():
268268
if handler.is_background:
269269
background += 1
270-
for field in state_cls.get_fields().values():
271-
key = _storage_feature_for_field(field)
270+
fields = state_cls.get_fields()
271+
for name in state_cls.base_vars:
272+
key = _storage_feature_for_field(fields[name])
272273
if key is not None:
273274
features[key] += 1
274275
features["shared_state_count"] = shared
Lines changed: 218 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,218 @@
1+
"""Integration test for the ``compile`` telemetry event.
2+
3+
Drives a real ``AppHarness`` compile against an app that exercises every
4+
field surfaced in ``features_used``, captures ``telemetry.send`` to a JSONL
5+
file inside the test app, then asserts the emitted counts.
6+
7+
The compile pipeline runs end-to-end here, so a regression in
8+
``record_compile`` (such as the inherited-storage double-count fix) shows
9+
up at the integration boundary rather than only in collector unit tests.
10+
"""
11+
12+
import functools
13+
import json
14+
from collections.abc import Generator
15+
from pathlib import Path
16+
17+
import pytest
18+
19+
from reflex.testing import AppHarness, chdir
20+
21+
22+
def TelemetryCompileApp(events_log_path: str = ""):
23+
"""Reflex app exercising the features collected on the ``compile`` event.
24+
25+
Includes a parent + child state hierarchy with Cookie / LocalStorage /
26+
SessionStorage to guard against the inherited-field double-count
27+
regression, a background event handler, a SharedState user subclass,
28+
and a dynamic route.
29+
30+
Args:
31+
events_log_path: Filesystem path to a JSONL file. Every
32+
``telemetry.send`` invocation during compile is appended as
33+
one JSON object per line so the test can read it back.
34+
"""
35+
import json
36+
import os
37+
from pathlib import Path
38+
39+
import reflex as rx
40+
from reflex.istate.storage import Cookie, LocalStorage, SessionStorage
41+
from reflex.utils import telemetry
42+
43+
# AppHarness force-disables telemetry. Re-enable it on the live config
44+
# singleton so the real _compile() path runs record_compile, and
45+
# redirect telemetry.send to disk so the test can read the payload.
46+
os.environ["REFLEX_TELEMETRY_ENABLED"] = "true"
47+
rx.config.get_config().telemetry_enabled = True
48+
49+
sink = Path(events_log_path)
50+
sink.parent.mkdir(parents=True, exist_ok=True)
51+
52+
def _capture(event, properties=None, **_kwargs):
53+
with sink.open("a") as fh:
54+
fh.write(
55+
json.dumps({"event": event, "properties": properties or {}}) + "\n"
56+
)
57+
return True
58+
59+
telemetry.send = _capture
60+
61+
class StorageRoot(rx.State):
62+
token: str = Cookie()
63+
local_pref: str = LocalStorage()
64+
65+
@rx.event(background=True)
66+
async def heavy(self):
67+
"""Background handler used to assert detection."""
68+
69+
class StorageChild(StorageRoot):
70+
# ``token`` and ``local_pref`` are inherited here and must not be
71+
# re-counted on the child state.
72+
session: str = SessionStorage()
73+
74+
class SharedThing(rx.SharedState):
75+
value: int = 0
76+
77+
def index():
78+
return rx.box(rx.text("home"))
79+
80+
def item_page():
81+
return rx.box(rx.text("item"))
82+
83+
app = rx.App()
84+
app.add_page(index)
85+
app.add_page(item_page, route="/items/[id]")
86+
87+
88+
@pytest.fixture
89+
def telemetry_events_log(tmp_path) -> Path:
90+
"""Path to the JSONL telemetry sink the harness app writes into.
91+
92+
Args:
93+
tmp_path: pytest tmp_path fixture.
94+
95+
Returns:
96+
Path the harness app writes one JSON event per line to.
97+
"""
98+
return tmp_path / "telemetry_events.jsonl"
99+
100+
101+
@pytest.fixture
102+
def telemetry_compile_harness(
103+
tmp_path_factory, telemetry_events_log: Path
104+
) -> Generator[AppHarness, None, None]:
105+
"""Run AppHarness for the telemetry-capture app.
106+
107+
Only ``_initialize_app()`` is invoked — that path imports the app
108+
module and calls ``App.__call__()`` which triggers ``_compile()``, so
109+
the captured JSONL contains the compile event by the time the test
110+
body runs. The dev backend / frontend lifecycle is skipped because
111+
this test only cares about compile-time telemetry.
112+
113+
Args:
114+
tmp_path_factory: pytest tmp_path_factory fixture.
115+
telemetry_events_log: Path the harness app writes events to.
116+
117+
Yields:
118+
The configured ``AppHarness`` instance after initialization.
119+
"""
120+
root = tmp_path_factory.mktemp("telemetry_compile_app")
121+
harness = AppHarness.create(
122+
root=root,
123+
app_source=functools.partial(
124+
TelemetryCompileApp,
125+
events_log_path=str(telemetry_events_log),
126+
),
127+
app_name="telemetry_compile_app",
128+
)
129+
harness._initialize_app()
130+
try:
131+
yield harness
132+
finally:
133+
if harness._registry_token is not None:
134+
from reflex_base.registry import RegistrationContext
135+
136+
RegistrationContext.reset(harness._registry_token)
137+
138+
139+
def _read_compile_events(events_log: Path) -> list[dict]:
140+
"""Return every ``compile`` event written to the sink.
141+
142+
Args:
143+
events_log: JSONL file the harness app appended to.
144+
145+
Returns:
146+
List of ``properties`` dicts for events named ``compile``.
147+
"""
148+
if not events_log.exists():
149+
return []
150+
out: list[dict] = []
151+
for line in events_log.read_text().splitlines():
152+
rec = json.loads(line)
153+
if rec.get("event") == "compile":
154+
out.append(rec["properties"])
155+
return out
156+
157+
158+
def test_compile_event_features_used_initial_and_hot_reload(
159+
telemetry_compile_harness: AppHarness,
160+
telemetry_events_log: Path,
161+
):
162+
"""Compile event payload is well-formed and stable across hot reload.
163+
164+
Two assertions in sequence, exercising the live ``_compile()`` path:
165+
166+
1. The initial compile (driven by AppHarness during fixture setup)
167+
emits a ``compile`` event with exact ``features_used`` counts.
168+
Guards the inherited client-storage double-count regression:
169+
``token`` and ``local_pref`` live on ``StorageRoot`` and are
170+
inherited by ``StorageChild`` — each must be counted exactly once.
171+
172+
2. A second ``_compile`` call under ``trigger="hot_reload"`` re-derives
173+
the snapshot from scratch and must produce the same counts. A
174+
regression that cached or accumulated counters across compiles
175+
shows up here as drift between the two events.
176+
"""
177+
payloads = _read_compile_events(telemetry_events_log)
178+
assert payloads, "no compile event was captured by the AppHarness run"
179+
initial = payloads[-1]
180+
initial_features = initial["features_used"]
181+
assert initial_features["cookie_count"] == 1, (
182+
"inherited cookie field was counted on parent and child"
183+
)
184+
assert initial_features["local_storage_count"] == 1, (
185+
"inherited LocalStorage field was counted on parent and child"
186+
)
187+
assert initial_features["session_storage_count"] == 1
188+
assert initial_features["background_event_handlers_count"] == 1
189+
assert initial_features["shared_state_count"] == 1
190+
assert initial_features["dynamic_routes_count"] == 1
191+
assert initial["trigger"] in {"initial", "backend_startup", None}
192+
assert initial["exception"] is None
193+
194+
app = telemetry_compile_harness.app_instance
195+
assert app is not None, "AppHarness did not populate app_instance"
196+
197+
# The real reflex CLI invokes hot reloads from inside the app directory.
198+
# AppHarness chdir's into app_path during _initialize_app() but reverts
199+
# on exit, so we restore it here to match the live hot-reload environment.
200+
pre_event_count = len(payloads)
201+
with chdir(telemetry_compile_harness.app_path):
202+
app._compile(trigger="hot_reload")
203+
204+
payloads = _read_compile_events(telemetry_events_log)
205+
assert len(payloads) == pre_event_count + 1, (
206+
"hot reload did not emit exactly one additional compile event"
207+
)
208+
reload = payloads[-1]
209+
assert reload["trigger"] == "hot_reload"
210+
assert reload["exception"] is None
211+
assert initial_features == reload["features_used"], (
212+
"features_used drifted between initial compile and hot reload"
213+
)
214+
assert initial["component_counts"] == reload["component_counts"]
215+
assert initial["pages_count"] == reload["pages_count"]
216+
assert [s["depth_from_root"] for s in initial["states"]] == [
217+
s["depth_from_root"] for s in reload["states"]
218+
]

tests/units/test_app.py

Lines changed: 106 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,7 @@
4545
from reflex.istate.manager.memory import StateManagerMemory
4646
from reflex.istate.manager.redis import StateManagerRedis
4747
from reflex.istate.manager.token import BaseStateToken
48+
from reflex.istate.storage import Cookie, LocalStorage, SessionStorage
4849
from reflex.model import Model
4950
from reflex.state import BaseState, OnLoadInternalState, State, reload_state_module
5051
from reflex.utils import exec as exec_utils
@@ -3065,6 +3066,111 @@ def test_compile_skips_telemetry_when_compile_app_short_circuits(
30653066
assert all(c.args[0] != "compile" for c in send_mock.call_args_list)
30663067

30673068

3069+
def test_compile_event_features_used_match_populated_app(
3070+
compilable_app: tuple[App, Path],
3071+
mocker: MockerFixture,
3072+
):
3073+
"""End-to-end compile fills features_used with exact counts from the app.
3074+
3075+
Guards against data-malformation regressions like the inherited-field
3076+
double-count: storage vars declared on a parent state must be counted
3077+
once, even when descendants inherit them.
3078+
"""
3079+
conf = rx.Config(app_name="testing", telemetry_enabled=True)
3080+
mocker.patch("reflex_base.config._get_config", return_value=conf)
3081+
app, web_dir = compilable_app
3082+
mocker.patch("reflex.utils.prerequisites.get_web_dir", return_value=web_dir)
3083+
mocker.patch("reflex.compiler.compiler.compile_app", return_value=True)
3084+
send_mock = mocker.patch("reflex.utils.telemetry.send")
3085+
3086+
class _PopulatedRoot(BaseState):
3087+
root_cookie: str = Cookie()
3088+
root_local: str = LocalStorage()
3089+
3090+
@rx.event(background=True)
3091+
async def heavy(self):
3092+
"""Background handler used to assert detection."""
3093+
3094+
class _PopulatedChild(_PopulatedRoot):
3095+
# root_cookie is inherited here and must not be re-counted.
3096+
child_session: str = SessionStorage()
3097+
3098+
app._state = _PopulatedRoot
3099+
3100+
def about_page():
3101+
return rx.box(rx.text("about"))
3102+
3103+
def item_page():
3104+
return rx.box(rx.text("item"))
3105+
3106+
app.add_page(about_page, route="/about")
3107+
app.add_page(item_page, route="/items/[id]")
3108+
3109+
app._compile(trigger="initial")
3110+
3111+
compile_calls = [c for c in send_mock.call_args_list if c.args[0] == "compile"]
3112+
assert len(compile_calls) == 1
3113+
payload = compile_calls[0].kwargs["properties"]
3114+
features = payload["features_used"]
3115+
assert features["cookie_count"] == 1, (
3116+
"inherited cookie field was counted on parent and child"
3117+
)
3118+
assert features["local_storage_count"] == 1
3119+
assert features["session_storage_count"] == 1
3120+
assert features["background_event_handlers_count"] == 1
3121+
assert features["dynamic_routes_count"] == 1
3122+
assert sorted(s["depth_from_root"] for s in payload["states"]) == [0, 1]
3123+
3124+
3125+
def test_compile_event_features_used_stable_across_hot_reload(
3126+
compilable_app: tuple[App, Path],
3127+
mocker: MockerFixture,
3128+
):
3129+
"""Recompiling under hot_reload reports the same features_used as the initial compile.
3130+
3131+
A bug that accumulated state across compiles (e.g. caching the snapshot
3132+
or re-counting on each pass) would show up here as drift between the
3133+
two payloads.
3134+
"""
3135+
conf = rx.Config(app_name="testing", telemetry_enabled=True)
3136+
mocker.patch("reflex_base.config._get_config", return_value=conf)
3137+
app, web_dir = compilable_app
3138+
mocker.patch("reflex.utils.prerequisites.get_web_dir", return_value=web_dir)
3139+
mocker.patch("reflex.compiler.compiler.compile_app", return_value=True)
3140+
send_mock = mocker.patch("reflex.utils.telemetry.send")
3141+
3142+
class _ReloadRoot(BaseState):
3143+
token: str = Cookie()
3144+
3145+
class _ReloadChild(_ReloadRoot):
3146+
session: str = SessionStorage()
3147+
3148+
app._state = _ReloadRoot
3149+
3150+
def page():
3151+
return rx.box(rx.text("hi"))
3152+
3153+
app.add_page(page, route="/items/[id]")
3154+
3155+
app._compile(trigger="initial")
3156+
app._compile(trigger="hot_reload")
3157+
3158+
compile_calls = [c for c in send_mock.call_args_list if c.args[0] == "compile"]
3159+
assert len(compile_calls) == 2
3160+
first = compile_calls[0].kwargs["properties"]
3161+
second = compile_calls[1].kwargs["properties"]
3162+
assert first["trigger"] == "initial"
3163+
assert second["trigger"] == "hot_reload"
3164+
assert first["features_used"] == second["features_used"]
3165+
assert [s["depth_from_root"] for s in first["states"]] == [
3166+
s["depth_from_root"] for s in second["states"]
3167+
]
3168+
# Sanity-check the values themselves, not just equality.
3169+
assert first["features_used"]["cookie_count"] == 1
3170+
assert first["features_used"]["session_storage_count"] == 1
3171+
assert first["features_used"]["dynamic_routes_count"] == 1
3172+
3173+
30683174
def test_call_marks_first_dev_backend_worker_as_startup(
30693175
compilable_app: tuple[App, Path],
30703176
mocker: MockerFixture,

tests/units/utils/test_telemetry_accounting.py

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -296,6 +296,31 @@ class _StorageState(BaseState):
296296
assert features["session_storage_count"] == 1
297297

298298

299+
def test_collect_features_used_storage_not_double_counted_through_inheritance():
300+
"""Inherited storage fields are not re-counted on each descendant."""
301+
302+
class _StorageParent(BaseState):
303+
c: str = Cookie()
304+
ls: str = LocalStorage()
305+
ss: str = SessionStorage()
306+
307+
class _StorageChild(_StorageParent):
308+
pass
309+
310+
class _StorageGrandchild(_StorageChild):
311+
extra: str = Cookie()
312+
313+
features = telemetry_accounting._collect_features_used(
314+
_fake_app(), # pyright: ignore[reportArgumentType]
315+
_fake_config(),
316+
[_StorageParent, _StorageChild, _StorageGrandchild],
317+
0,
318+
)
319+
assert features["cookie_count"] == 2
320+
assert features["local_storage_count"] == 1
321+
assert features["session_storage_count"] == 1
322+
323+
299324
def test_collect_features_used_upload_from_component_walk():
300325
"""``upload_count`` is read from the component walk, not a marker."""
301326
features = telemetry_accounting._collect_features_used(

0 commit comments

Comments
 (0)