Add basic websockets support (#1027)

Author: wwwillchen

.github/workflows/ci.yml

@@ -72,6 +72,14 @@ jobs:
           name: playwright-report-with-concurrent-updates-enabled
           path: playwright-report-with-concurrent-updates-enabled/
           retention-days: 30
+      - name: Run playwright test with websockets enabled
+        run: MESOP_WEBSOCKETS_ENABLED=true PLAYWRIGHT_HTML_OUTPUT_DIR=playwright-report-with-websockets-enabled yarn playwright test
+      - uses: actions/upload-artifact@a8a3f3ad30e3422c9c7b888a15615d19a852ae32 # v3.1.3
+        if: always()
+        with:
+          name: playwright-report-with-websockets-enabled
+          path: playwright-report-with-websockets-enabled/
+          retention-days: 30
       - name: Run playwright test with memory state session
         run: MESOP_STATE_SESSION_BACKEND=memory PLAYWRIGHT_HTML_OUTPUT_DIR=playwright-report-with-memory-state-session yarn playwright test
       - uses: actions/upload-artifact@a8a3f3ad30e3422c9c7b888a15615d19a852ae32 # v3.1.3

.vscode/settings.json

@@ -23,5 +23,6 @@
     "bazel-out/**": true,
     "bazel-mesop/**": true
   },
-  "python.analysis.extraPaths": ["./bazel-bin"]
+  "python.analysis.extraPaths": ["./bazel-bin"],
+  "typescript.tsdk": "node_modules/typescript/lib"
 }

build_defs/defaults.bzl

@@ -77,6 +77,10 @@ THIRD_PARTY_JS_HIGHLIGHTJS = [
     "@npm//highlight.js",
 ]
 
+THIRD_PARTY_JS_SOCKETIO_CLIENT = [
+    "@npm//socket.io-client",
+]
+
 THIRD_PARTY_PY_ABSL_PY = [
     requirement("absl-py"),
 ]
@@ -89,6 +93,10 @@ THIRD_PARTY_PY_FLASK = [
     requirement("flask"),
 ]
 
+THIRD_PARTY_PY_FLASK_SOCKETIO = [
+    requirement("flask-socketio"),
+]
+
 THIRD_PARTY_PY_MATPLOTLIB = [
     requirement("matplotlib"),
 ]

build_defs/requirements.txt

@@ -8,6 +8,8 @@ python-dotenv
 
 # Optional (lazily-loaded) deps:
 sqlalchemy
+flask-socketio
+
 # greenlet is needed for SQL Alchemy depending on the architecture, but because of how
 # Bazel works using requirements_lock.txt, it does seem to able to install the
 # architecture specific requirements (in this case caught on Github CI).

build_defs/requirements_lock.txt

@@ -16,6 +16,10 @@ babel==2.15.0 \
     --hash=sha256:08706bdad8d0a3413266ab61bd6c34d0c28d6e1e7badf40a2cebe67644e2e1fb \
     --hash=sha256:8daf0e265d05768bc6c7a314cf1321e9a123afc328cc635c18622a2f30a04413
     # via mkdocs-material
+bidict==0.23.1 \
+    --hash=sha256:03069d763bc387bbd20e7d49914e75fc4132a41937fa3405417e1a5a2d006d71 \
+    --hash=sha256:5dae8d4d79b552a71cbabc7deb25dfe8ce710b17ff41711e13010ead2abfc3e5
+    # via python-socketio
 blinker==1.8.2 \
     --hash=sha256:1779309f71bf239144b9399d06ae925637cf6634cf6bd131104184531bf67c01 \
     --hash=sha256:8f77b09d3bf7c795e969e9486f39c2c5e9c39d4ee07424be2bc594ece9642d83
@@ -290,6 +294,12 @@ firebase-admin==6.5.0 \
 flask==3.0.3 \
     --hash=sha256:34e815dfaa43340d1d15a5c3a02b8476004037eb4840b34910c6e21679d288f3 \
     --hash=sha256:ceb27b0af3823ea2737928a4d99d125a06175b8512c445cbd9a9ce200ef76842
+    # via
+    #   -r build_defs/requirements.txt
+    #   flask-socketio
+flask-socketio==5.4.1 \
+    --hash=sha256:2e9b8864a5be37ca54f6c76a4d06b1ac5e0df61fde12d03afc81ab4057e1eb86 \
+    --hash=sha256:895da879d162781b9193cbb8fe8f3cf25b263ff242980d5c5e6c16d3c03930d2
     # via -r build_defs/requirements.txt
 fonttools==4.53.0 \
     --hash=sha256:099634631b9dd271d4a835d2b2a9e042ccc94ecdf7e2dd9f7f34f7daf333358d \
@@ -580,6 +590,10 @@ grpcio-status==1.62.2 \
     --hash=sha256:206ddf0eb36bc99b033f03b2c8e95d319f0044defae9b41ae21408e7e0cda48f \
     --hash=sha256:62e1bfcb02025a1cd73732a2d33672d3e9d0df4d21c12c51e0bbcaf09bab742a
     # via google-api-core
+h11==0.14.0 \
+    --hash=sha256:8f19fbbe99e72420ff35c00b27a34cb9937e902a8b810e2c88300c6f0a3b699d \
+    --hash=sha256:e3fe4ac4b851c468cc8363d500db52c2ead036020723024a109d37346efaa761
+    # via wsproto
 httplib2==0.22.0 \
     --hash=sha256:14ae0a53c1ba8f3d37e9e27cf37eabb0fb9980f435ba405d546948b009dd64dc \
     --hash=sha256:d7a10bc5ef5ab08322488bde8c726eeee5c8618723fdb399597ec58f3d82df81
@@ -716,7 +730,6 @@ markdown==3.6 \
     --hash=sha256:48f276f4d8cfb8ce6527c8f79e2ee29708508bf4d40aa410fbc3b4ee832c850f \
     --hash=sha256:ed4f41f6daecbeeb96e576ce414c41d2d876daa9a16cb35fa8ed8c2ddfad0224
     # via
-    #   -r build_defs/requirements.txt
     #   mkdocs
     #   mkdocs-autorefs
     #   mkdocs-material
@@ -1231,9 +1244,7 @@ pydantic-core==2.18.4 \
 pygments==2.18.0 \
     --hash=sha256:786ff802f32e91311bff3889f6e9a86e81505fe99f2735bb6d60ae0c5004f199 \
     --hash=sha256:b8e6aca0523f3ab76fee51799c488e38782ac06eafcf95e7ba832985c8e7b13a
-    # via
-    #   -r build_defs/requirements.txt
-    #   mkdocs-material
+    # via mkdocs-material
 pyjwt[crypto]==2.8.0 \
     --hash=sha256:57e28d156e3d5c10088e0c68abb90bfac3df82b40a71bd0daa20c65ccd5c23de \
     --hash=sha256:59127c392cc44c2da5bb3192169a91f429924e17aff6534d70fdc02ab3e04320
@@ -1265,6 +1276,14 @@ python-dotenv==1.0.1 \
     --hash=sha256:e324ee90a023d808f1959c46bcbc04446a10ced277783dc6ee09987c37ec10ca \
     --hash=sha256:f7b63ef50f1b690dddf550d03497b66d609393b40b564ed0d674909a68ebf16a
     # via -r build_defs/requirements.txt
+python-engineio==4.9.1 \
+    --hash=sha256:7631cf5563086076611e494c643b3fa93dd3a854634b5488be0bba0ef9b99709 \
+    --hash=sha256:f995e702b21f6b9ebde4e2000cd2ad0112ba0e5116ec8d22fe3515e76ba9dddd
+    # via python-socketio
+python-socketio==5.11.4 \
+    --hash=sha256:42efaa3e3e0b166fc72a527488a13caaac2cefc76174252486503bd496284945 \
+    --hash=sha256:8b0b8ff2964b2957c865835e936310190639c00310a47d77321a594d1665355e
+    # via flask-socketio
 pytz==2024.1 \
     --hash=sha256:2a29735ea9c18baf14b448846bde5a48030ed267578472d8955cd0e7443a9812 \
     --hash=sha256:328171f4e3623139da4983451950b28e95ac706e13f3f2630a879749e7a8b319
@@ -1423,6 +1442,10 @@ rsa==4.9 \
     --hash=sha256:90260d9058e514786967344d0ef75fa8727eed8a7d2e43ce9f4bcf1b536174f7 \
     --hash=sha256:e38464a49c6c85d7f1351b0126661487a7e0a14a50f1675ec50eb34d4f20ef21
     # via google-auth
+simple-websocket==1.1.0 \
+    --hash=sha256:4af6069630a38ed6c561010f0e11a5bc0d4ca569b36306eb257cd9a192497c8c \
+    --hash=sha256:7939234e7aa067c534abdab3a9ed933ec9ce4691b0713c78acb195560aa52ae4
+    # via python-engineio
 six==1.16.0 \
     --hash=sha256:1e61c37477a1626458e36f7b1d82aa5c9b094fa4802892072e49de9c60c4c926 \
     --hash=sha256:8abb2f1d86890a2dfb989f9a77cfcfd3e47c2a354b01111771326f8aa26e0254
@@ -1543,3 +1566,7 @@ werkzeug==3.0.3 \
     --hash=sha256:097e5bfda9f0aba8da6b8545146def481d06aa7d3266e7448e2cccf67dd8bd18 \
     --hash=sha256:fc9645dc43e03e4d630d23143a04a7f947a9a3b5727cd535fdfe155a17cc48c8
     # via flask
+wsproto==1.2.0 \
+    --hash=sha256:ad565f26ecb92588a3e43bc3d96164de84cd9902482b130d0ddbaa9664a85065 \
+    --hash=sha256:b9acddd652b585d75b20477888c56642fdade28bdfd3579aa24a4d2c037dd736
+    # via simple-websocket

mesop/cli/cli.py

@@ -22,6 +22,7 @@
 from mesop.server.flags import port
 from mesop.server.logging import log_startup
 from mesop.server.server import configure_flask_app
+from mesop.server.server_utils import MESOP_WEBSOCKETS_ENABLED
 from mesop.server.static_file_serving import configure_static_file_serving
 from mesop.utils.host_util import get_public_host
 from mesop.utils.runfiles import get_runfile_location
@@ -153,7 +154,17 @@ def main(argv: Sequence[str]):
     log_startup(port=port())
     logging.getLogger("werkzeug").setLevel(logging.WARN)
 
-  flask_app.run(host=get_public_host(), port=port(), use_reloader=False)
+  if MESOP_WEBSOCKETS_ENABLED:
+    socketio = flask_app.socketio  # type: ignore
+    socketio.run(
+      flask_app,
+      host=get_public_host(),
+      port=port(),
+      use_reloader=False,
+      allow_unsafe_werkzeug=True,
+    )
+  else:
+    flask_app.run(host=get_public_host(), port=port(), use_reloader=False)
 
 
 if __name__ == "__main__":

mesop/components/input/e2e/input_test.ts

@@ -19,6 +19,7 @@ test('test input on_blur works', async ({page}) => {
   // Fill in input and then click button and make sure values match
   await page.getByLabel('Input').click();
   await page.getByLabel('Input').fill('abc');
+  await page.getByLabel('Input').blur();
   await page.getByRole('button', {name: 'button'}).click();
   await expect(page.getByText('Input: abc')).toBeVisible();
   await expect(
@@ -28,6 +29,7 @@ test('test input on_blur works', async ({page}) => {
   // Same with textarea:
   await page.getByLabel('Regular textarea').click();
   await page.getByLabel('Regular textarea').fill('123');
+  await page.getByLabel('Regular textarea').blur();
   await page.getByRole('button', {name: 'button'}).click();
   await expect(page.getByText('Input: 123')).toBeVisible();
   await expect(
@@ -37,6 +39,7 @@ test('test input on_blur works', async ({page}) => {
   // Same with native textarea:
   await page.getByRole('textbox').nth(2).click();
   await page.getByRole('textbox').nth(2).fill('second_textarea');
+  await page.getByRole('textbox').nth(2).blur();
   await page.getByRole('button', {name: 'button'}).click();
   await expect(page.getByText('Input: second_textarea')).toBeVisible();
   await expect(

mesop/examples/e2e/chat_test.ts

@@ -8,6 +8,7 @@ test('Chat UI can send messages and display responses', async ({page}) => {
 
   // Test that we can send a message.
   await page.locator('//input').fill('Lorem ipsum');
+  await page.locator('//input').blur();
   // Need to wait for the input state to be saved before clicking.
   await page.waitForTimeout(2000);
   await page.getByRole('button').filter({hasText: 'send'}).click();

mesop/protos/ui.proto

@@ -179,6 +179,7 @@ message RenderEvent {
 message ExperimentSettings {
     optional bool experimental_editor_toolbar_enabled = 1;
     optional bool concurrent_updates_enabled = 2;
+    optional bool websockets_enabled = 3;
 }
 
 // UI response event for updating state.

mesop/server/BUILD

@@ -4,6 +4,7 @@ load(
     "THIRD_PARTY_PY_DOTENV",
     "THIRD_PARTY_PY_FIREBASE_ADMIN",
     "THIRD_PARTY_PY_FLASK",
+    "THIRD_PARTY_PY_FLASK_SOCKETIO",
     "THIRD_PARTY_PY_GREENLET",
     "THIRD_PARTY_PY_MSGPACK",
     "THIRD_PARTY_PY_PYTEST",
@@ -37,16 +38,18 @@ py_library(
                "//mesop/utils",
                "//mesop/warn",
            ] + THIRD_PARTY_PY_ABSL_PY +
-           THIRD_PARTY_PY_FLASK,
+           THIRD_PARTY_PY_FLASK +
+           THIRD_PARTY_PY_FLASK_SOCKETIO,
 )
 
 py_library(
     name = "state_sessions",
     srcs = STATE_SESSIONS_SRCS,
     deps = [
-        "//mesop/dataclass_utils",
-        "//mesop/exceptions",
-    ] + THIRD_PARTY_PY_MSGPACK + THIRD_PARTY_PY_DOTENV,
+               "//mesop/dataclass_utils",
+               "//mesop/exceptions",
+           ] + THIRD_PARTY_PY_MSGPACK +
+           THIRD_PARTY_PY_DOTENV,
 )
 
 py_test(