Update app.py

app.py

@@ -309,6 +309,12 @@ class GradioUI:
             return message, "error"
 
     def interact_with_agent(self, prompt, messages, session_state):
+        # Clear any stale session data at the beginning
+        if "hf_token" in session_state and not session_state.get("hf_token"):
+            del session_state["hf_token"]
+        if "agent" in session_state and not session_state.get("agent"):
+            del session_state["agent"]
+        
         # Get or create session-specific agent
         if "agent" not in session_state:
             # Check if we have a valid HF_TOKEN in session
@@ -433,6 +439,16 @@ User Query: """
                 
                 messages.append(msg)
                 yield messages
+            
+            # Clear sensitive data from session after interaction (AUTOMATIC)
+            if "hf_token" in session_state:
+                del session_state["hf_token"]
+            if "agent" in session_state:
+                del session_state["agent"]
+            if "HF_TOKEN" in os.environ:
+                del os.environ["HF_TOKEN"]
+            logger.info("Session automatically cleared after interaction")
+            
             yield messages
         except Exception as e:
             logger.error(f"Error in interaction: {str(e)}")
@@ -461,7 +477,7 @@ User Query: """
         is_valid, message = validate_hf_api_key(token_to_use)
         
         if is_valid:
-            # Store HF_TOKEN in session state
+            # Store HF_TOKEN in session state (but will be cleared after use)
             session_state["hf_token"] = token_to_use
             session_state["max_steps"] = max_steps
             logger.info(f"API key stored in session from {source}: {token_to_use[:10]}...")
@@ -483,6 +499,8 @@ User Query: """
             logger.warning(f"Invalid API key from {source}: {token_to_use[:10] if token_to_use else 'None'}...")
             return f"❌ Invalid API key from {source}: {message}"
 
+
+
     def upload_file(
         self,
         file,
@@ -675,7 +693,7 @@ This AI agent specializes in automated vulnerability research and analysis, buil
                             # API Key Configuration Section
                             with gr.Accordion("🔑 API Configuration", open=False):
                                 gr.Markdown("**Configure your Hugging Face API Key**")
-                                gr.Markdown("All API keys are stored only in session memory, not persisted.")
+                                gr.Markdown("🔒 **Security**: API keys are automatically cleared after each interaction for your privacy.")
                                 gr.Markdown("Get your API key from: https://huggingface.co/settings/tokens")
                                 
                                 api_key_input = gr.Textbox(
@@ -750,6 +768,8 @@ This AI agent specializes in automated vulnerability research and analysis, buil
                             [api_key_input, max_steps_slider, session_state],
                             [api_key_status]
                         )
+                        
+
 
                         text_input.submit(
                             self.log_user_message,
@@ -826,7 +846,7 @@ This AI agent specializes in automated vulnerability research and analysis, buil
                         # API Key Configuration Section for Mobile
                         with gr.Accordion("🔑 API Configuration", open=False):
                             gr.Markdown("**Configure your Hugging Face API Key**")
-                            gr.Markdown("Due to recent API changes, you need to provide your own Hugging Face API key to use this application.")
+                            gr.Markdown("🔒 **Security**: API keys are automatically cleared after each interaction for your privacy.")
                             gr.Markdown("Get your API key from: https://huggingface.co/settings/tokens")
                             
                             mobile_api_key_input = gr.Textbox(
@@ -878,6 +898,8 @@ This AI agent specializes in automated vulnerability research and analysis, buil
                             [mobile_api_key_status]
                         )
                         
+
+                        
                         # Mobile Example button events
                         mobile_example_btn_1 = gr.Button("🔍 MobaXterm 24.0 vulnerabilities", size="sm", variant="secondary")
                         mobile_example_btn_2 = gr.Button("🔍 Chrome 120.0.6099.109 security analysis", size="sm", variant="secondary")