Fix Tavily API key persistence in session state

When users update the Tavily API key through the UI on HF Spaces, the agent now properly stores the key in session_state and recreates with both the OAuth token and the Tavily API key, ensuring MCP search tools are initialized correctly.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <noreply@anthropic.com>

my_ui.py

@@ -22,19 +22,23 @@ class CustomGradioUI(GradioUI):
         self.allowed_file_types = allowed_file_types or [".pdf", ".docx", ".txt"]
         self.examples = examples or []
 
-    def update_api_key(self, api_key, current_key, oauth_token: gr.OAuthToken | None = None):
+    def update_api_key(self, api_key, current_key, session_state, oauth_token: gr.OAuthToken | None = None):
         """Update the agent with a new API key."""
         if api_key and api_key != current_key:
             # Recreate the agent with the new API key
             self.agent = self.agent_factory(tavily_api_key=api_key, oauth_token=oauth_token)
-            return api_key, gr.Markdown("✓ API key updated successfully", visible=True)
+            # Store in session state for persistence
+            session_state["tavily_api_key"] = api_key
+            return api_key, gr.Markdown("✓ API key updated successfully", visible=True), session_state
         elif not api_key and current_key:
             # Reset to default (env var)
             self.agent = self.agent_factory(oauth_token=oauth_token)
+            # Clear from session state
+            session_state.pop("tavily_api_key", None)
             return "", gr.Markdown(
                 "API key cleared, using environment variable if set", visible=True
-            )
-        return current_key, gr.Markdown("", visible=False)
+            ), session_state
+        return current_key, gr.Markdown("", visible=False), session_state
 
     def interact_with_agent_oauth(self, stored_messages, chatbot, session_state, oauth_token: gr.OAuthToken | None = None):
         """Wrapper for interact_with_agent that recreates agent with OAuth token."""
@@ -117,8 +121,8 @@ class CustomGradioUI(GradioUI):
                     # Update agent when API key changes
                     tavily_api_key_input.change(
                         self.update_api_key,
-                        [tavily_api_key_input, current_api_key],
-                        [current_api_key, api_key_status],
+                        [tavily_api_key_input, current_api_key, session_state],
+                        [current_api_key, api_key_status, session_state],
                     )
 
                 gr.HTML(