Added code snapshot download

app.py

@@ -206,6 +206,19 @@ def initialize_state():
     st.session_state.final_deployment_path = None
     st.session_state.final_deployment_pdf_path = None
     # --- END ADDED ---
+    # Snapshot paths
+    st.session_state.snapshot_path_codegen_initial = None
+    st.session_state.snapshot_path_codegen_refined = None
+    st.session_state.snapshot_path_review_refined = None
+    st.session_state.snapshot_path_testing_refined = None
+    st.session_state.snapshot_path_qa_polished = None
+    # Keep final folder paths (might point to last snapshot or dedicated folder)
+    st.session_state.review_code_snapshot_folder = None # Points to post_review snapshot
+    st.session_state.testing_passed_code_folder = None # Points to snapshot saved by save_testing_outputs
+    st.session_state.final_code_folder = None # Points to post_qa snapshot
+    # --- END ADDED/MODIFIED ---
+
+    logger.info("Streamlit session state initialized/reset including PDF and Snapshot paths.")
 
     logger.info("Streamlit session state initialized/reset including PDF paths.")
 
@@ -495,32 +508,68 @@ with st.sidebar:
         # Indicates UML stage hasn't run or failed before saving
         st.caption("*UML diagrams have not been generated yet.*")
 
-    # Code Snapshot Downloads (ZIP)
+ # --- MODIFIED: Code Snapshot Downloads (ZIP) ---
     st.markdown("---")
     st.subheader("Code Snapshots (ZIP)")
-    st.caption("Download code versions from key stages.")
-    # Use the helper function for ZIP creation and download
+    st.caption("Download code versions from various stages.")
+
+    # Code Generation Cycle Snapshots
+    st.markdown("**Code Generation Cycle:**")
     create_zip_and_download_button(
-        folder_path_key="review_code_snapshot_folder",
-        zip_path_key="review_code_zip_path",
-        zip_basename="code_snapshot_review",
-        button_label_prefix="Review Stage Code",
+        folder_path_key="snapshot_path_codegen_initial", # Use the new state key
+        zip_path_key="zip_path_cg_initial", # Unique key for session state zip path
+        zip_basename="snapshot_codegen_initial",
+        button_label_prefix="Initial Code",
         sidebar_context=st.sidebar
     )
     create_zip_and_download_button(
-        folder_path_key="testing_passed_code_folder",
-        zip_path_key="testing_code_zip_path",
-        zip_basename="code_snapshot_testing",
-        button_label_prefix="Testing Stage Code",
+        folder_path_key="snapshot_path_codegen_refined", # Use the new state key
+        zip_path_key="zip_path_cg_refined", # Unique key
+        zip_basename="snapshot_codegen_refined",
+        button_label_prefix="Refined Code (Latest)", # Label indicates latest
         sidebar_context=st.sidebar
     )
+    st.markdown("---")
+
+    # Review & Security Cycle Snapshot
+    st.markdown("**Review & Security Cycle:**")
+    create_zip_and_download_button(
+        folder_path_key="snapshot_path_review_refined", # Use the new state key
+        zip_path_key="zip_path_review_refined", # Unique key
+        zip_basename="snapshot_review_refined",
+        button_label_prefix="Post-Review Code",
+        sidebar_context=st.sidebar
+    )
+    st.markdown("---")
+
+    # Testing Cycle Snapshots
+    st.markdown("**Testing Cycle:**")
+    create_zip_and_download_button(
+        folder_path_key="snapshot_path_testing_refined", # Use the new state key for failed refinement
+        zip_path_key="zip_path_testing_refined", # Unique key
+        zip_basename="snapshot_testing_refined",
+        button_label_prefix="Post-Failure Refined Code (Latest)",
+        sidebar_context=st.sidebar
+    )
+    create_zip_and_download_button(
+        folder_path_key="testing_passed_code_folder", # Keep the one saved on PASS
+        zip_path_key="zip_path_testing_passed", # Unique key
+        zip_basename="snapshot_testing_passed",
+        button_label_prefix="Passed Testing Code",
+        sidebar_context=st.sidebar
+    )
+    st.markdown("---")
+
+    # Quality Analysis Cycle Snapshot (Final Code)
+    st.markdown("**Quality Analysis Cycle:**")
     create_zip_and_download_button(
-        folder_path_key="final_code_folder",
-        zip_path_key="final_code_zip_path",
-        zip_basename="code_snapshot_final",
-        button_label_prefix="Final Code",
+        folder_path_key="snapshot_path_qa_polished", # Use the new state key (points to final code)
+        zip_path_key="zip_path_qa_polished", # Unique key
+        zip_basename="snapshot_qa_polished_final",
+        button_label_prefix="Final Polished Code",
         sidebar_context=st.sidebar
     )
+    # --- END MODIFIED Code Snapshots ---
 
     st.divider()
 
@@ -853,6 +902,12 @@ with main_col:
                                     "final_deployment_path": None, # MD path
                                     "final_deployment_pdf_path": None, # PDF path
                                     # --- END ADDED/MODIFIED ---
+                                    # Intermediate Snapshot Paths
+                                    "snapshot_path_codegen_initial": None,
+                                    "snapshot_path_codegen_refined": None,
+                                    "snapshot_path_review_refined": None,
+                                    "snapshot_path_testing_refined": None,
+                                    "snapshot_path_qa_polished": None,
                                 }
 
                                 # Update the main session state variables