From 8fabc6cb7d2547cff1571a331b7aa6be61252d4a Mon Sep 17 00:00:00 2001
From: Will Chen <willchen90@gmail.com>
Date: Tue, 3 Dec 2024 15:28:52 -0800
Subject: [PATCH] Skip unnecessary trace render loop in WebSockets mode (#1123)

---
 mesop/runtime/context.py |  7 +------
 mesop/runtime/runtime.py |  4 +---
 mesop/server/server.py   | 10 +++++++---
 3 files changed, 9 insertions(+), 12 deletions(-)

diff --git a/mesop/runtime/context.py b/mesop/runtime/context.py
index 323cddfd..0e4de245 100644
--- a/mesop/runtime/context.py
+++ b/mesop/runtime/context.py
@@ -35,7 +35,6 @@ def __init__(
     self._states: dict[type[Any], object] = states
     # Previous states is used for performing state diffs.
     self._previous_states: dict[type[Any], object] = copy.deepcopy(states)
-    self._trace_mode = False
     self._handlers: dict[str, Handler] = {}
     self._commands: list[pb.Command] = []
     self._node_slot: pb.Component | None = None
@@ -191,11 +190,7 @@ def viewport_size(self) -> pb.ViewportSize:
     return self._viewport_size
 
   def register_event_handler(self, fn_id: str, handler: Handler) -> None:
-    if self._trace_mode:
-      self._handlers[fn_id] = handler
-
-  def set_trace_mode(self, trace_mode: bool) -> None:
-    self._trace_mode = trace_mode
+    self._handlers[fn_id] = handler
 
   def current_node(self) -> pb.Component:
     return self._current_node
diff --git a/mesop/runtime/runtime.py b/mesop/runtime/runtime.py
index f17571ef..3fa5ae0e 100644
--- a/mesop/runtime/runtime.py
+++ b/mesop/runtime/runtime.py
@@ -108,9 +108,7 @@ def wait_for_hot_reload(self):
         lambda: not self.is_hot_reload_in_progress, initial_delay=0.100
       )
 
-  def run_path(self, path: str, trace_mode: bool = False) -> None:
-    self.context().set_trace_mode(trace_mode)
-
+  def run_path(self, path: str) -> None:
     if path not in self._path_to_page_config:
       paths = list(self._path_to_page_config.keys())
       if not paths:
diff --git a/mesop/server/server.py b/mesop/server/server.py
index 68ad51b8..2c4f0d91 100644
--- a/mesop/server/server.py
+++ b/mesop/server/server.py
@@ -76,7 +76,7 @@ def render_loop(
   ) -> Generator[str, None, None]:
     try:
       runtime().context().acquire_lock()
-      runtime().run_path(path=path, trace_mode=trace_mode)
+      runtime().run_path(path=path)
       page_config = runtime().get_page_config(path=path)
       title = page_config.title if page_config else "Unknown path"
 
@@ -196,8 +196,12 @@ def generate_data(ui_request: pb.UiRequest) -> Generator[str, None, None]:
           else:
             runtime().context().restore_state_from_session(event.state_token)
 
-        for _ in render_loop(path=ui_request.path, trace_mode=True):
-          pass
+        # 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:
+          for _ in render_loop(path=ui_request.path, trace_mode=True):
+            pass
         if ui_request.user_event.handler_id:
           runtime().context().set_previous_node_from_current_node()
         else: