Update app.py via AI Editor

app.py

@@ -1,5 +1,5 @@
 import dash
-from dash import dcc, html, Input, Output, State, callback, ALL
+from dash import dcc, html, Input, Output, State, ALL, callback_context
 import dash_bootstrap_components as dbc
 import base64
 import io
@@ -25,11 +25,13 @@ SESSION_DATA = {}
 SESSION_LOCKS = {}
 
 def get_session_id(set_cookie=False):
-    # Try to get session id from browser cookies
-    sid = request.cookies.get('session-id')
+    sid = None
+    try:
+        sid = request.cookies.get('session-id')
+    except Exception:
+        pass
     if sid and sid in SESSION_DATA:
         return sid
-    # If not present, generate a new one
     sid = str(uuid.uuid4())
     if set_cookie:
         g.set_cookie_sid = sid
@@ -48,6 +50,17 @@ def get_session_data():
         SESSION_LOCKS[sid] = threading.Lock()
     return SESSION_DATA[sid], SESSION_LOCKS[sid]
 
+def restore_session_files(session_data):
+    temp_dir = session_data.get('temp_dir')
+    if not temp_dir or not os.path.exists(temp_dir):
+        return
+    files = os.listdir(temp_dir)
+    for filename in files:
+        file_path = os.path.join(temp_dir, filename)
+        if filename not in session_data['uploaded_files']:
+            session_data['uploaded_files'][filename] = file_path
+            session_data['file_texts'][filename] = parse_file_content(file_path, filename)
+
 def cleanup_session_tempdirs():
     for sess in SESSION_DATA.values():
         try:
@@ -82,8 +95,11 @@ openai.api_key = os.environ.get('OPENAI_API_KEY')
 
 @server.before_request
 def ensure_session_cookie():
-    # Set session cookie if not present
-    sid = request.cookies.get('session-id')
+    sid = None
+    try:
+        sid = request.cookies.get('session-id')
+    except Exception:
+        sid = None
     if not sid or sid not in SESSION_DATA:
         sid_new = str(uuid.uuid4())
         g.set_cookie_sid = sid_new
@@ -279,15 +295,16 @@ def update_output(list_of_contents, list_of_names):
     raise PreventUpdate
 
 @app.callback(
-    Output('file-list', 'children', allow_duplicate=True),
+    Output('file-list', 'children'),
     Input({'type': 'remove-file', 'index': ALL}, 'n_clicks'),
+    State('file-list', 'children'),
     prevent_initial_call=True
 )
-def remove_file(n_clicks):
-    ctx = dash.callback_context
-    session_data, lock = get_session_data()
-    if not ctx.triggered:
+def remove_file(n_clicks, file_list_children):
+    ctx = callback_context
+    if not ctx.triggered or not any(n_clicks):
         raise PreventUpdate
+    session_data, lock = get_session_data()
     triggered_id = ctx.triggered[0]['prop_id'].split('.')[0]
     import ast
     try:
@@ -299,6 +316,7 @@ def remove_file(n_clicks):
         if removed_file in session_data['uploaded_files']:
             try:
                 os.remove(session_data['uploaded_files'][removed_file])
+                logger.info(f"Deleted file from disk: {removed_file}")
             except Exception as e:
                 logger.warning(f"Failed to delete temp file {removed_file}: {e}")
             session_data['uploaded_files'].pop(removed_file, None)
@@ -306,6 +324,17 @@ def remove_file(n_clicks):
         logger.info(f"Files after deletion: {list(session_data['uploaded_files'].keys())}")
     return get_file_cards(session_data['uploaded_files'])
 
+@app.callback(
+    Output('file-list', 'children', allow_duplicate=True),
+    Input('file-list', 'children'),
+    prevent_initial_call=False
+)
+def restore_file_list(file_list_children):
+    session_data, lock = get_session_data()
+    with lock:
+        restore_session_files(session_data)
+        return get_file_cards(session_data['uploaded_files'])
+
 @app.callback(
     Output('matrix-preview', 'children'),
     Output('loading-output', 'children'),
@@ -314,7 +343,7 @@ def remove_file(n_clicks):
 )
 def generate_matrix(*args):
     session_data, lock = get_session_data()
-    ctx = dash.callback_context
+    ctx = callback_context
     if not ctx.triggered:
         raise PreventUpdate
     button_id = ctx.triggered[0]['prop_id'].split('.')[0]