fix generate_data to do trace render loop for web sockets (#1345)

nirben82 · web-flow · commit a030ae62ca50 · 2026-02-19T13:02:07.000-08:00
diff --git a/mesop/runtime/context.py b/mesop/runtime/context.py
@@ -116,6 +116,7 @@ def __init__(
     self._previous_states: dict[type[Any], object] = copy.deepcopy(states)
     self._handlers: dict[str, Handler] = {}
     self._commands: list[pb.Command] = []
+    self._has_rendered: bool = False
     self._viewport_size: pb.ViewportSize | None = None
     self._theme_settings: pb.ThemeSettings | None = None
     self._js_modules: set[str] = set()
@@ -146,6 +147,12 @@ def js_modules(self) -> set[str]:
   def clear_js_modules(self):
     self._js_modules = set()
 
+  def has_rendered(self) -> bool:
+    return self._has_rendered
+
+  def set_has_rendered(self, has_rendered: bool) -> None:
+    self._has_rendered = has_rendered
+
   def query_params(self) -> dict[str, list[str]]:
     return self._query_params
 
diff --git a/mesop/server/server.py b/mesop/server/server.py
@@ -130,6 +130,7 @@ def render_loop(
           ],
         )
       )
+      runtime().context().set_has_rendered(True)
       yield serialize(data)
     except Exception as e:
       logging.error(e)
@@ -205,16 +206,20 @@ def generate_data(ui_request: pb.UiRequest) -> Generator[str, None, None]:
         runtime().context().set_viewport_size(event.viewport_size)
         runtime().context().initialize_query_params(event.query_params)
 
-        if not MESOP_WEBSOCKETS_ENABLED:
+        if (
+          not MESOP_WEBSOCKETS_ENABLED or not runtime().context().has_rendered()
+        ):
           if event.states.states:
             runtime().context().update_state(event.states)
           else:
             runtime().context().restore_state_from_session(event.state_token)
 
-        # In websockets mode, we don't need to do a trace render loop because
-        # the context instance is long-lived and contains all the registered
-        # event handlers from the last render loop.
-        if not MESOP_WEBSOCKETS_ENABLED:
+        # In websockets mode, since the context instance is long-lived, we only
+        # need to do a trace render loop if the context has not completed at least
+        # one render loop.
+        if (
+          not MESOP_WEBSOCKETS_ENABLED or not runtime().context().has_rendered()
+        ):
           for _ in render_loop(path=ui_request.path, trace_mode=True):
             pass
         if ui_request.user_event.handler_id:
