changes

app.py

@@ -34,6 +34,10 @@ if IS_HUGGINGFACE:
     app.config["SESSION_COOKIE_SECURE"] = True
     app.config["SESSION_COOKIE_HTTPONLY"] = True
     app.config["SESSION_COOKIE_SAMESITE"] = "Lax"
+    # Add more durable storage for Hugging Face
+    app.config["SESSION_FILE_THRESHOLD"] = 10  # Low threshold to ensure writes
+    # Important: Don't use large session lifetime on Hugging Face
+    app.config["PERMANENT_SESSION_LIFETIME"] = 60 * 60 * 24  # 1 day only
 
 Session(app)
 
@@ -44,6 +48,8 @@ MODEL_MAPPING = {
     "gemini-2.0-flash-thinking": "gemini-2.0-flash-thinking-exp-01-21"
 }
 
+global_api_key = None
+
 def get_api_model_name(frontend_model_name):
     return MODEL_MAPPING.get(frontend_model_name, frontend_model_name)
 
@@ -56,18 +62,15 @@ def set_api_key():
     api_key = request.form.get("api_key")
     if not api_key:
         return jsonify({"success": False, "message": "API key is required"}), 400
-    
     try:
         # Test the API key
         genai.configure(api_key=api_key)
         model = genai.GenerativeModel("gemini-2.0-pro-exp-02-05")
         response = model.generate_content("Say 'API key is valid'")
-        
         # Store API key in session with permanent flag
         session.permanent = True
         session["api_key"] = api_key
         logger.info("API key successfully set and validated")
-        
         return jsonify({"success": True})
     except Exception as e:
         logger.error(f"API key validation error: {str(e)}")
@@ -78,7 +81,6 @@ def generate_notebook_route():
     if "api_key" not in session:
         logger.warning("Generate notebook request without API key")
         return jsonify({"success": False, "message": "API key not set"}), 401
-    
     # Handle both GET (for streaming) and POST requests
     if request.method == "GET":
         prompt = request.args.get("prompt")
@@ -100,14 +102,12 @@ def generate_notebook_route():
     
     try:
         genai.configure(api_key=session["api_key"])
-        
         # OPTIMIZATION: If format_only is True, skip the AI call and just format the provided content
         if request.method == "POST" and format_only:
             # Use client-provided content as is (it's already the AI response)
             notebook_content = prompt
             notebook_json = format_notebook(notebook_content)
             notebook_info = extract_notebook_info(notebook_content)
-            
             return jsonify({
                 "success": True,
                 "notebook": notebook_json,
@@ -121,7 +121,6 @@ def generate_notebook_route():
             notebook_content = generate_notebook(prompt, api_model_name)
             notebook_json = format_notebook(notebook_content)
             notebook_info = extract_notebook_info(notebook_content)
-            
             return jsonify({
                 "success": True,
                 "notebook": notebook_json,
@@ -177,9 +176,8 @@ def edit_notebook_route():
     # Map the frontend model name to the API model name
     api_model_name = get_api_model_name(model_name)
     
-    genai.configure(api_key=session["api_key"])
-    
     try:
+        genai.configure(api_key=session["api_key"])
         if stream:
             return stream_notebook_edit(edit_prompt, notebook_json, api_model_name)
         else:
@@ -187,7 +185,6 @@ def edit_notebook_route():
             edited_content = edit_notebook(edit_prompt, notebook_json, api_model_name)
             notebook_json = format_notebook(edited_content)
             notebook_info = extract_notebook_info(edited_content)
-            
             return jsonify({
                 "success": True,
                 "notebook": notebook_json,
@@ -195,13 +192,11 @@ def edit_notebook_route():
                 "description": notebook_info["description"]
             })
     except Exception as e:
-        app.logger.error(f"Error editing notebook: {str(e)}")
+        logger.error(f"Error editing notebook: {str(e)}", exc_info=True)
         return jsonify({"success": False, "message": str(e)}), 500
 
 @app.route("/download_notebook", methods=["POST"])
 def download_notebook():
-    from flask import Response
-    
     data = request.json
     notebook_json = data.get("notebook")
     filename = data.get("filename", f"notebook_{uuid.uuid4()}.ipynb")
@@ -217,10 +212,43 @@ def download_notebook():
         mimetype="application/json",
         headers={"Content-Disposition": f"attachment;filename={filename}"}
     )
-    
     return response
 
+@app.route("/session_test", methods=["GET"])
+def session_test():
+    if "session_test" not in session:
+        session["session_test"] = True
+        is_new = True
+    else:
+        is_new = False
+    
+    # Log the current session state
+    session_vars = list(session.keys()) if session else []
+    logger.info(f"Session check - variables: {session_vars}")
+    
+    # Check if API key is available from URL parameter (fallback for HF)
+    api_key_param = request.args.get('api_key_param')
+    if api_key_param and "api_key" not in session:
+        global global_api_key
+        try:
+            # Validate it quickly before accepting
+            genai.configure(api_key=api_key_param)
+            model = genai.GenerativeModel("gemini-2.0-pro-exp-02-05")
+            # If no exception, store it
+            session["api_key"] = api_key_param
+            global_api_key = api_key_param
+            logger.info("API key set from URL parameter")
+        except Exception as e:
+            logger.error(f"Invalid API key from URL parameter: {str(e)}")
+    
+    return jsonify({
+        "session_works": True,
+        "is_new_session": is_new,
+        "has_api_key": "api_key" in session,
+        "session_vars": session_vars
+    })
+
 if __name__ == "__main__":
-    port = int(os.environ.get("PORT", 5000))
     debug_mode = os.environ.get("FLASK_ENV") == "development"
-    app.run(host="0.0.0.0", port=port, debug=debug_mode)
+    port = int(os.environ.get("PORT", 5000))
+    app.run(host="0.0.0.0", port=port, debug=debug_mode)

static/js/main.js

@@ -166,6 +166,9 @@ document.addEventListener('DOMContentLoaded', function() {
                 // Store API key in localStorage as backup
                 localStorage.setItem('notegenie_api_key', apiKey);
                 
+                // For Hugging Face, add a flag that API key was set in this session
+                sessionStorage.setItem('api_key_set', 'true');
+                
                 if (document.querySelector('#apiKeyModal.show')) {
                     showApiKeyFeedback('API key saved successfully!', 'success');
                     setTimeout(() => {
@@ -302,12 +305,24 @@ document.addEventListener('DOMContentLoaded', function() {
         function createEventSource() {
             // Create a new event source with a unique timestamp to prevent caching
             const timestamp = Date.now();
+            
+            // On Hugging Face, add API key as URL parameter if stored (fallback mechanism)
+            let apiKeyParam = '';
+            if (window.location.hostname.includes('huggingface.co') || 
+                window.location.hostname.includes('hf.space')) {
+                const storedApiKey = localStorage.getItem('notegenie_api_key');
+                if (storedApiKey && !sessionStorage.getItem('api_key_tried_in_url')) {
+                    apiKeyParam = `&api_key_param=${encodeURIComponent(storedApiKey)}`;
+                    sessionStorage.setItem('api_key_tried_in_url', 'true');
+                }
+            }
+            
             eventSource = new EventSource(`/generate_notebook?${new URLSearchParams({
                 prompt: prompt,
                 model: modelName,
                 stream: true,
                 t: timestamp  // Add timestamp to prevent caching
-            }).toString()}`);
+            }).toString()}${apiKeyParam}`);
             
             eventSource.onmessage = function(event) {
                 // Update our last-activity timestamp
@@ -408,7 +423,17 @@ document.addEventListener('DOMContentLoaded', function() {
                     
                     // Check if it's an auth error (most likely API key not set)
                     if (err.status === 401) {
-                        updateAiMessage(aiMessageId, '**Error: API key not set or invalid.** \n\nPlease click the API Key button in the top right corner to set your Google Gemini API key.');
+                        // For Hugging Face: If we have an API key in storage but it's not being recognized
+                        // in the session, show a special message
+                        if ((window.location.hostname.includes('huggingface.co') || 
+                             window.location.hostname.includes('hf.space')) && 
+                            localStorage.getItem('notegenie_api_key') && 
+                            sessionStorage.getItem('api_key_set')) {
+                            
+                            updateAiMessage(aiMessageId, '**Error: Session issue with API key.** \n\nPlease try refreshing the page completely, then set your API key again.');
+                        } else {
+                            updateAiMessage(aiMessageId, '**Error: API key not set or invalid.** \n\nPlease click the API Key button in the top right corner to set your Google Gemini API key.');
+                        }
                         showApiKeyModal();
                     } else {
                         // Try to salvage what we have so far

templates/index.html

@@ -229,5 +229,28 @@
     <script src="https://cdnjs.cloudflare.com/ajax/libs/prism/1.29.0/plugins/autoloader/prism-autoloader.min.js"></script>
     <script src="https://cdn.jsdelivr.net/npm/marked/marked.min.js"></script>
     <script src="{{ url_for('static', filename='js/main.js') }}"></script>
+    
+    <!-- Hugging Face Spaces Session Check -->
+    <script>
+    if (window.location.hostname.includes('huggingface.co') || window.location.hostname.includes('hf.space')) {
+        console.log("Running on Hugging Face Spaces");
+        // Check if cookies are enabled
+        if (navigator.cookieEnabled) {
+            console.log("Cookies are enabled");
+        } else {
+            console.warn("Cookies are disabled - sessions won't work!");
+        }
+        
+        // Add a hidden diagonstic endpoint
+        fetch('/session_check')
+            .then(response => response.json())
+            .then(data => {
+                console.log("Session check:", data);
+            })
+            .catch(err => {
+                console.error("Session check failed:", err);
+            });
+    }
+    </script>
 </body>
 </html>

utils/ai_helpers.py

@@ -115,8 +115,12 @@ def edit_notebook(edit_request, notebook_json, model_name="gemini-2.0-pro-exp-02
     
     return response.text
 
-def stream_notebook_generation(user_prompt, model_name="gemini-2.0-pro-exp-02-05"):
+def stream_notebook_generation(user_prompt, model_name="gemini-2.0-pro-exp-02-05", api_key=None):
     """Stream notebook generation responses from Gemini API."""
+    # Configure the API with the key if provided
+    if api_key:
+        genai.configure(api_key=api_key)
+        
     model = genai.GenerativeModel(model_name)
     enhanced_prompt = craft_notebook_prompt(user_prompt)
     
@@ -208,8 +212,12 @@ def stream_notebook_generation(user_prompt, model_name="gemini-2.0-pro-exp-02-05
         headers=headers
     )
 
-def stream_notebook_edit(edit_request, notebook_json, model_name="gemini-2.0-pro-exp-02-05"):
+def stream_notebook_edit(edit_request, notebook_json, model_name="gemini-2.0-pro-exp-02-05", api_key=None):
     """Stream notebook editing responses from Gemini API."""
+    # Configure the API with the key if provided
+    if api_key:
+        genai.configure(api_key=api_key)
+        
     model = genai.GenerativeModel(model_name)
     enhanced_prompt = craft_edit_prompt(edit_request, notebook_json)