Update main.py

main.py

@@ -558,7 +558,6 @@ def generate_sozo_report(project_id):
         if not uid:
             return jsonify({'error': 'Unauthorized'}), 401
         
-        # Check user credits first
         user_ref = db.reference(f'users/{uid}')
         user_data = user_ref.get()
         if not user_data:
@@ -574,13 +573,11 @@ def generate_sozo_report(project_id):
         if not project_data or project_data.get('uid') != uid:
             return jsonify({'error': 'Project not found or unauthorized'}), 404
         
-        # --- THE AUTHORITATIVE LOCK ---
         current_status = project_data.get('status')
-        if current_status in ['generating_report', 'generating_video']:
+        if current_status in ['generating_report', 'generating_video', 'generating_slides']:
             logger.warning(f"User {uid} attempted to generate a report for project {project_id} which is already in progress (status: {current_status}).")
-            return jsonify({'error': 'Report generation is already in progress for this project.'}), 409 # 409 Conflict
+            return jsonify({'error': 'A process is already running for this project.'}), 409
 
-        # Set the lock BEFORE starting the heavy work
         project_ref.update({'status': 'generating_report'})
         logger.info(f"Project {project_id} status locked to 'generating_report'.")
         
@@ -597,29 +594,44 @@ def generate_sozo_report(project_id):
             bucket
         )
         
-        update_data = {
-            'status': 'draft', # Unlock by setting the new status
-            'rawMarkdown': draft_data.get('raw_md'),
-            'chartUrls': draft_data.get('chartUrls'),
-            'dataContext': draft_data.get('data_context')
-        }
-        
-        project_ref.update(update_data)
-        
-        # Deduct credits ONLY after successful generation
-        user_ref.update({'credits': current_credits - 2})
-        
-        logger.info(f"Project {project_id} successfully updated with draft data. User {uid} charged 2 credits.")
-        
-        return jsonify({
-            'success': True,
-            'project': {**project_data, **update_data},
-            'credits_remaining': current_credits - 2
-        }), 200
-        
+        # --- Graceful Fallback Save & Charge Logic ---
+        try:
+            # Happy Path: Try to save everything
+            full_update_data = {
+                'status': 'draft',
+                'rawMarkdown': draft_data.get('raw_md'),
+                'chartUrls': draft_data.get('chartUrls'),
+                'dataContext': draft_data.get('data_context')
+            }
+            project_ref.update(full_update_data)
+            user_ref.update({'credits': current_credits - 2})
+            logger.info(f"Project {project_id} successfully updated with full data. User {uid} charged 2 credits.")
+            return jsonify({
+                'success': True,
+                'project': {**project_data, **full_update_data},
+                'credits_remaining': current_credits - 2
+            }), 200
+        except Exception as save_error:
+            # Fallback Path: Save only the essentials if the full save fails
+            logger.warning(f"Failed to save full project data for {project_id} due to: {save_error}. Saving degraded version.")
+            degraded_update_data = {
+                'status': 'draft',
+                'rawMarkdown': draft_data.get('raw_md'),
+                'chartUrls': draft_data.get('chartUrls'),
+                'dataContext': None,
+                'warning': 'Report generated, but context data failed to save. Downstream features may be affected.'
+            }
+            project_ref.update(degraded_update_data)
+            user_ref.update({'credits': current_credits - 2})
+            logger.info(f"Project {project_id} successfully updated with DEGRADED data. User {uid} charged 2 credits.")
+            return jsonify({
+                'success': True,
+                'project': {**project_data, **degraded_update_data},
+                'credits_remaining': current_credits - 2
+            }), 200
+
     except Exception as e:
         logger.error(f"CRITICAL error generating report for project {project_id}: {traceback.format_exc()}")
-        # Make sure to unlock the project on failure
         db.reference(f'sozo_projects/{project_id}').update({
             'status': 'failed',
             'error': str(e)