Update app.py via AI Editor

app.py

@@ -153,7 +153,6 @@ 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(
             [
                 dbc.Col(
@@ -203,17 +202,39 @@ app.layout = dbc.Container(
     className="p-4"
 )
 
+@app.server.before_request
+def persist_session_cookie():
+    session_id = flask.request.cookies.get('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
+        # Attach the response only if needed
+        flask.request.session_id_set = session_id
+    else:
+        flask.g.session_id = session_id
+
 @app.callback(
     Output('session-id-store', 'data'),
-    Input('dummy-div', 'children'),
+    Input('session-id-store', 'data'),
     prevent_initial_call=False
 )
-def ensure_session_id(_):
-    sid = get_session_id_from_cookie()
-    if not sid:
+def ensure_session_id(session_id):
+    # On first load, set session-id-store from cookie or generate new
+    try:
+        sid = session_id
+        if not sid:
+            if hasattr(flask.request, 'session_id_set'):
+                sid = flask.request.session_id_set
+            else:
+                sid = get_or_create_session_id()
+        flask.g.session_id = sid
+        return sid
+    except Exception as e:
+        logging.error(f"Error ensuring session id: {e}")
         sid = str(uuid.uuid4())
-    flask.g.session_id = sid
-    return sid
+        return sid
 
 @app.callback(
     Output('file-info', 'children'),
@@ -232,7 +253,6 @@ def ensure_session_id(_):
 def handle_upload(contents, filename, clear_n, delete_upload_n_list, split_n, session_data, session_id):
     trigger = ctx.triggered_id
 
-    # Always use session id from store
     if not session_id:
         session_id = str(uuid.uuid4())
     flask.g.session_id = session_id
@@ -397,26 +417,6 @@ def download_zip_file(session_id, filename):
         logging.error(f"ZIP file not found for download: {file_path}")
         return "File not found", 404
 
-@app.callback(
-    Output('dummy-div', 'children'),
-    Input('session-id-store', 'data'),
-    prevent_initial_call=False
-)
-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_from_cookie()
-    if not session_id:
-        session_id = str(uuid.uuid4())
-    flask.g.session_id = session_id
-
 if __name__ == '__main__':
     print("Starting the Dash application...")
     app.run(debug=True, host='0.0.0.0', port=7860, threaded=True)