Fix state management issues with Close Document button and Process Document Again button

app.py

@@ -527,6 +527,10 @@ if 'temp_file_paths' not in st.session_state:
 if 'last_processed_file' not in st.session_state:
     st.session_state.last_processed_file = None
 
+# Ensure perform_reset flag is initialized
+if 'perform_reset' not in st.session_state:
+    st.session_state.perform_reset = False
+
 # Check if we need to perform a complete reset (coming from "Close Document" button)
 if 'perform_reset' in st.session_state and st.session_state.perform_reset:
     # List of all session state keys that should be reset, except previous_results
@@ -1388,37 +1392,40 @@ with main_tab1:
         with left_col:
             # Process button styling is now handled by global CSS
             
-            # Make the button more clear about its function
-            if st.session_state.processed_document_active:
-                process_button = st.button("Process Document Again")
+            # Use a key for the button based on state to force re-creation
+            button_key = "process_again" if st.session_state.processed_document_active else "process_initial"
+            
+            # Show appropriate button text based on state
+            button_text = "Process Document Again" if st.session_state.processed_document_active else "Process Document"
+            
+            # Create the button
+            process_button = st.button(button_text, key=button_key)
+            
+            # Handle sample document recreation if needed
+            if process_button and st.session_state.processed_document_active and st.session_state.original_sample_bytes is not None:
+                # Recreate the uploaded file from stored bytes
+                from io import BytesIO
+                import mimetypes
                 
-                # If process button is clicked and we're dealing with a sample document that was cleared
-                if process_button and uploaded_file is None and 'original_sample_bytes' in st.session_state:
-                    # Recreate the uploaded file from stored bytes
-                    from io import BytesIO
-                    import mimetypes
-                    
-                    # Determine mime type based on file extension
-                    file_ext = os.path.splitext(st.session_state.original_sample_name)[1].lower()
-                    if file_ext == '.pdf':
-                        mime_type = 'application/pdf'
-                    elif file_ext in ['.jpg', '.jpeg']:
-                        mime_type = 'image/jpeg'
-                    elif file_ext == '.png':
-                        mime_type = 'image/png'
-                    else:
-                        mime_type = mimetypes.guess_type(st.session_state.original_sample_name)[0] or 'application/octet-stream'
-                    
-                    # Create a synthetic file-like object with the same interface as UploadedFile
-                    uploaded_file = type('obj', (object,), {
-                        'name': st.session_state.original_sample_name,
-                        'getvalue': lambda: st.session_state.original_sample_bytes,
-                        'read': lambda: st.session_state.original_sample_bytes,
-                        'seek': lambda x: None,
-                        'type': mime_type
-                    })
-            else:
-                process_button = st.button("Process Document")
+                # Determine mime type based on file extension
+                file_ext = os.path.splitext(st.session_state.original_sample_name)[1].lower()
+                if file_ext == '.pdf':
+                    mime_type = 'application/pdf'
+                elif file_ext in ['.jpg', '.jpeg']:
+                    mime_type = 'image/jpeg'
+                elif file_ext == '.png':
+                    mime_type = 'image/png'
+                else:
+                    mime_type = mimetypes.guess_type(st.session_state.original_sample_name)[0] or 'application/octet-stream'
+                
+                # Create a synthetic file-like object with the same interface as UploadedFile
+                uploaded_file = type('obj', (object,), {
+                    'name': st.session_state.original_sample_name,
+                    'getvalue': lambda: st.session_state.original_sample_bytes,
+                    'read': lambda: st.session_state.original_sample_bytes,
+                    'seek': lambda x: None,
+                    'type': mime_type
+                })
             
             # Empty container for progress indicators - will be filled during processing
             # Positioned right after the process button for better visibility
@@ -2256,9 +2263,12 @@ with main_tab1:
                 with success_cols[1]:
                     # Close button styling is now handled by global CSS
                     
-                    if st.button("✕ Close Document", key="close_document_button", help="Clear current document and start over"):
-                        # Create a special flag to signal a complete reset on the next rerun
-                        st.session_state.perform_reset = True
+                    # Define a function to clear document state
+                    def clear_document_state():
+                        # Reset all document-related session state
+                        st.session_state.processed_document_active = False
+                        st.session_state.sample_document = None
+                        st.session_state.last_processed_file = None
                         
                         # Clean up any temporary files
                         if 'temp_file_paths' in st.session_state:
@@ -2268,8 +2278,14 @@ with main_tab1:
                                         os.remove(temp_path)
                                 except Exception:
                                     pass  # Ignore errors in cleanup
-                        
-                        st.rerun()
+                            # Clear the temp files list
+                            st.session_state.temp_file_paths = []
+                    
+                    # Create the close button with a callback
+                    st.button("✕ Close Document", 
+                              key="close_document_button", 
+                              help="Clear current document and start over",
+                              on_click=clear_document_state)
                     
                 # Store the result in the previous results list
                 # Add timestamp to result for history tracking