Update app.py via AI Editor

app.py

@@ -19,7 +19,11 @@ logging.basicConfig(level=logging.INFO, format='%(asctime)s %(levelname)s %(mess
 SESSION_DATA = {}
 SESSION_LOCKS = {}
 
-def get_session_id():
+def get_session_id_from_cookie():
+    sid = flask.request.cookies.get('session-id')
+    return sid
+
+def get_or_create_session_id():
     sid = flask.request.cookies.get('session-id')
     if not sid:
         sid = str(uuid.uuid4())
@@ -147,6 +151,7 @@ def get_split_results_placeholder():
 
 app.layout = dbc.Container(
     [
+        dcc.Store(id='session-id-store', storage_type='session'),
         dcc.Store(id='session-store', storage_type='session'),
         html.Div(id='dummy-div', style={'display': 'none'}),
         dbc.Row(
@@ -198,6 +203,18 @@ app.layout = dbc.Container(
     className="p-4"
 )
 
+@app.callback(
+    Output('session-id-store', 'data'),
+    Input('dummy-div', 'children'),
+    prevent_initial_call=False
+)
+def ensure_session_id(_):
+    sid = get_session_id_from_cookie()
+    if not sid:
+        sid = str(uuid.uuid4())
+    flask.g.session_id = sid
+    return sid
+
 @app.callback(
     Output('file-info', 'children'),
     Output('split-btn', 'disabled'),
@@ -209,12 +226,13 @@ app.layout = dbc.Container(
     Input({'type': 'delete-upload-btn', 'index': ALL}, 'n_clicks'),
     Input('split-btn', 'n_clicks'),
     State('session-store', 'data'),
+    State('session-id-store', 'data'),
     prevent_initial_call=True
 )
-def handle_upload(contents, filename, clear_n, delete_upload_n_list, split_n, session_data):
+def handle_upload(contents, filename, clear_n, delete_upload_n_list, split_n, session_data, session_id):
     trigger = ctx.triggered_id
-    # Always get current session id from cookie
-    session_id = flask.request.cookies.get('session-id')
+
+    # Always use session id from store
     if not session_id:
         session_id = str(uuid.uuid4())
     flask.g.session_id = session_id
@@ -381,18 +399,22 @@ def download_zip_file(session_id, filename):
 
 @app.callback(
     Output('dummy-div', 'children'),
-    Input('session-store', 'data'),
-    prevent_initial_call=True
+    Input('session-id-store', 'data'),
+    prevent_initial_call=False
 )
-def set_cookie_on_load(session_data):
-    session_id = get_session_id()
+def set_cookie_on_load(session_id):
+    if not session_id:
+        session_id = str(uuid.uuid4())
     resp = flask.make_response("")
     resp.set_cookie('session-id', session_id, max_age=60*60*24*3)
+    flask.g.session_id = session_id
     return ""
 
 @app.server.before_request
 def persist_session_cookie():
-    session_id = get_session_id()
+    session_id = get_session_id_from_cookie()
+    if not session_id:
+        session_id = str(uuid.uuid4())
     flask.g.session_id = session_id
 
 if __name__ == '__main__':