From cc3c7d5d3b46b4f1ffb5ce068f5f2c236c2df177 Mon Sep 17 00:00:00 2001 From: Will Chen Date: Tue, 29 Oct 2024 14:43:33 -0700 Subject: [PATCH] Move copy_current_request_context to ensure thread-safety --- mesop/server/server.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/mesop/server/server.py b/mesop/server/server.py index b2280f79..16d017b4 100644 --- a/mesop/server/server.py +++ b/mesop/server/server.py @@ -292,7 +292,6 @@ def teardown_clear_stale_state_sessions(error=None): @sock.route(UI_PATH) def handle_websocket(ws: Server): - @copy_current_request_context def ws_generate_data(ws, ui_request): for data_chunk in generate_data(ui_request): if not ws.connected: @@ -319,8 +318,13 @@ def ws_generate_data(ws, ui_request): # Start a new thread so we can handle multiple # concurrent updates for the same websocket connection. + # + # Note: we do copy_current_request_context at the callsite + # to ensure that the request context is copied over for each new thread. thread = threading.Thread( - target=ws_generate_data, args=(ws, ui_request), daemon=True + target=copy_current_request_context(ws_generate_data), + args=(ws, ui_request), + daemon=True, ) thread.start()