minor tweaks

mesop/server/server.py

@@ -24,13 +24,11 @@
 from mesop.utils.url_utils import remove_url_query_param
 from mesop.warn import warn
 
+UI_PATH = "/__ui__"
+
 
 def configure_flask_app(
-  *,
-  prod_mode: bool = True,
-  exceptions_to_propagate: Sequence[type] = (),
-  # TODO: plumb this from an env var
-  is_websockets_enabled=True,
+  *, prod_mode: bool = True, exceptions_to_propagate: Sequence[type] = ()
 ) -> Flask:
   flask_app = Flask(__name__)
 
@@ -224,7 +222,7 @@ def generate_data(ui_request: pb.UiRequest) -> Generator[str, None, None]:
         error=pb.ServerError(exception=str(e), traceback=format_traceback())
       )
 
-  @flask_app.route("/__ui__", methods=["POST"])
+  @flask_app.route(UI_PATH, methods=["POST"])
   def ui_stream() -> Response:
     # Prevent CSRF by checking the request site matches the site
     # of the URL root (where the Flask app is being served from)
@@ -234,7 +232,7 @@ def ui_stream() -> Response:
     if not runtime().debug_mode and not is_same_site(
       request.headers.get("Origin"), request.url_root
     ):
-      abort(403, "Rejecting cross-site POST request to /__ui__")
+      abort(403, "Rejecting cross-site POST request to " + UI_PATH)
     data = request.data
     if not data:
       raise Exception("Missing request payload")
@@ -251,12 +249,12 @@ def teardown_clear_stale_state_sessions(error=None):
   if not prod_mode:
     configure_debug_routes(flask_app)
 
-  if is_websockets_enabled:
+  if MESOP_WEBSOCKETS_ENABLED:
     from flask_socketio import SocketIO, emit
 
     socketio = SocketIO(flask_app)
 
-    @socketio.on("message", namespace="/__ui__")
+    @socketio.on("message", namespace=UI_PATH)
     def handle_message(message):
       if not message:
         emit("error", {"error": "Missing request payload"})

mesop/server/wsgi_app.py

@@ -31,7 +31,10 @@ def run(self):
         use_reloader=False,
         allow_unsafe_werkzeug=True,
       )
-    self._flask_app.run(host=get_local_host(), port=port(), use_reloader=False)
+    else:
+      self._flask_app.run(
+        host=get_local_host(), port=port(), use_reloader=False
+      )
 
 
 def create_app(

mesop/web/src/services/channel.ts

@@ -22,6 +22,8 @@ import {getQueryParams} from '../utils/query_params';
 import {ExperimentService} from './experiment_service';
 import {io, Socket} from 'socket.io-client'; // Import Socket.IO client
 
+const STREAM_END = '<stream_end>';
+
 // Pick 500ms as the minimum duration before showing a progress/busy indicator
 // for the channel.
 // See: https://github.com/google/mesop/issues/365
@@ -137,17 +139,14 @@ export class Channel {
       // Looks like Angular has a bug where it's not intercepting EventSource onmessage.
       zone.run(() => {
         const data = (e as any).data;
-        if (data === '<stream_end>') {
+        if (data === STREAM_END) {
           this.eventSource.close();
           this.status = ChannelStatus.CLOSED;
           clearTimeout(this.isWaitingTimeout);
           this.isWaiting = false;
           this._isHotReloading = false;
           this.logger.log({type: 'StreamEnd'});
-          if (this.queuedEvents.length) {
-            const queuedEvent = this.queuedEvents.shift()!;
-            queuedEvent();
-          }
+          this.dequeueEvent();
           return;
         }
 
@@ -159,9 +158,6 @@ export class Channel {
     });
   }
 
-  /**
-   * Initialize WebSocket connection using Socket.IO.
-   */
   private initWebSocket(initParams: InitParams, request: UiRequest) {
     if (this.socket) {
       this.status = ChannelStatus.OPEN;
@@ -194,14 +190,11 @@ export class Channel {
       const prefix = 'data: ';
       const payloadData = (data.data.slice(prefix.length) as string).trimEnd();
       zone.run(() => {
-        if (payloadData === '<stream_end>') {
+        if (payloadData === STREAM_END) {
           this._isHotReloading = false;
           this.status = ChannelStatus.CLOSED;
           this.logger.log({type: 'StreamEnd'});
-          if (this.queuedEvents.length) {
-            const queuedEvent = this.queuedEvents.shift()!;
-            queuedEvent();
-          }
+          this.dequeueEvent();
           return;
         }
 
@@ -229,9 +222,6 @@ export class Channel {
     });
   }
 
-  /**
-   * Handle UiResponse from the server.
-   */
   private handleUiResponse(
     request: UiRequest,
     uiResponse: UiResponse,
@@ -408,6 +398,13 @@ export class Channel {
     }
   }
 
+  private dequeueEvent() {
+    if (this.queuedEvents.length) {
+      const queuedEvent = this.queuedEvents.shift()!;
+      queuedEvent();
+    }
+  }
+
   checkForHotReload() {
     const pollHotReloadEndpoint = async () => {
       try {