Update app.py

app.py

@@ -157,8 +157,11 @@ def configure_gemini(api_key: str | None, model_name: str | None) -> str:
 def call_gemini(prompt: str, api_key: str, model_name: str, use_grounding: bool = False) -> str:
     """Calls the Gemini API with a given prompt, optionally using grounding."""
     # This check is crucial - it will raise an error *before* the API call if prereqs aren't met
-    if not api_key or not model_name: # Check for empty string "" as well as None
-        raise ValueError("Gemini API key or model not set.")
+    # Check for empty string "" as well as None
+    if not isinstance(api_key, str) or api_key == "" or not model_name:
+        # This error indicates a failure in the workflow logic or state propagation
+        # because this function should only be called when prereqs are met.
+        raise ValueError(f"Gemini API call prerequisites not met: api_key={api_key}, model_name={model_name}")
     try:
         genai.configure(api_key=api_key)
         model = genai.GenerativeModel(model_name)
@@ -236,9 +239,9 @@ def check_send_button_ready(profile: gr.OAuthProfile | None, token: gr.OAuthToke
     print(f"  Received token type: {type(token)}, is None: {token is None}")
     # For api_key, print part of the key if not None for verification, be careful with full key
     # Also check for empty string for printing
-    api_key_display = api_key[:5] if api_key else ('Empty String' if isinstance(api_key, str) and api_key == "" else 'N/A')
-    print(f"  Received api_key is None: {api_key is None}, is Empty String: {api_key == '' if isinstance(api_key, str) else 'N/A'}, first 5 chars or status: {api_key_display}")
-    print(f"  Received model_name: {model_name}")
+    api_key_display = api_key[:5] if isinstance(api_key, str) and api_key else ('Empty String' if isinstance(api_key, str) and api_key == "" else 'None')
+    print(f"  Received api_key: Type={type(api_key)}, Value={api_key_display}")
+    print(f"  Received model_name: Type={type(model_name)}, Value={model_name}")
     # --- END ENHANCED DEBUGGING LOGS ---
 
     is_logged_in    = profile is not None and token is not None
@@ -265,11 +268,9 @@ def ai_workflow_chat(
     history: list[dict],
     hf_profile: gr.OAuthProfile | None,
     hf_token:   gr.OAuthToken   | None,
-    # gemini_api_key and gemini_model are inputs, but their *state* is managed by change handlers,
-    # so we don't yield them back in this version. They are read from the state at the start
-    # of the function call.
-    gemini_api_key: str         | None, # Read key from state
-    gemini_model:   str         | None, # Read model from state
+    # Pass gemini_api_key and gemini_model as inputs - these come from the State variables
+    gemini_api_key: str         | None,
+    gemini_model:   str         | None,
     repo_id_state:  str | None,
     workflow_state: str,
     space_sdk:      str,
@@ -297,7 +298,8 @@ def ai_workflow_chat(
     str | None,       # 8: Updated repo name (for repo_name_state state)
     str | None,       # 9: Updated generated code (for generated_code_state state)
     bool,             # 10: Updated use_grounding_state (for use_grounding_state state)
-    # gemini_key, gemini_model removed from outputs based on new key handling strategy
+    str | None,       # 11: Explicitly yield gemini_api_key State
+    str | None,       # 12: Explicitly yield gemini_model State
 ]:
     """
     Generator function to handle the AI workflow state machine.
@@ -328,9 +330,10 @@ def ai_workflow_chat(
     # Yield immediately to update the chat UI with the user's message
     # This provides immediate feedback to the user while the AI processes
     # Ensure all state variables and UI outputs are yielded back in the correct order
-    # NOTE: gemini_key and gemini_model are NOT yielded back
+    # Include gemini_api_key and gemini_model in the yield tuple
     yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-           attempts, app_desc, repo_name, generated_code, use_grounding)
+           attempts, app_desc, repo_name, generated_code, use_grounding,
+           gemini_api_key, gemini_model) # Explicitly pass back current state values
 
     try:
         # --- State Machine Logic based on the current 'state' variable ---
@@ -339,17 +342,17 @@ def ai_workflow_chat(
         # the checks remain as a safeguard for the workflow logic itself.
         if not (hf_profile and hf_token):
             history = add_bot_message(history, "Workflow paused: Please log in to Hugging Face first.")
-            # NOTE: gemini_key and gemini_model are NOT yielded back
             yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                   attempts, app_desc, repo_name, generated_code, use_grounding)
+                   attempts, app_desc, repo_name, generated_code, use_grounding,
+                   gemini_api_key, gemini_model) # Explicitly pass back current state values
             return # Stop workflow execution for this click
 
         # Check if API key is non-empty string and model is set
         if not (isinstance(gemini_api_key, str) and gemini_api_key != "" and gemini_model):
              history = add_bot_message(history, "Workflow cannot start: Please ensure your Gemini API key is entered and a model is selected.")
-             # NOTE: gemini_key and gemini_model are NOT yielded back
              yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                    attempts, app_desc, repo_name, generated_code, use_grounding)
+                    attempts, app_desc, repo_name, generated_code, use_grounding,
+                    gemini_api_key, gemini_model) # Explicitly pass back current state values
              return # Stop workflow execution for this click
 
 
@@ -366,9 +369,10 @@ def ai_workflow_chat(
                 history = add_bot_message(history, "Workflow reset.")
                 # Yield updated history and reset state variables to their initial values
                 # Also reset UI outputs to their initial state
-                # NOTE: gemini_key and gemini_model are NOT yielded back (they are managed by change handlers)
+                # gemini_key and gemini_model are passed back with their *current* (likely valid) state
                 yield (history, None, STATE_IDLE, "<p>No Space created yet.</p>", "", "", 0,
-                       None, None, None, False) # Reset use_grounding to default False, and other states to None/default
+                       None, None, None, False,
+                       gemini_api_key, gemini_model) # Reset use_grounding to default False, other states to None/default
                 # No return needed after yield in this generator pattern; execution for this click ends here.
 
             elif generate_match:
@@ -381,9 +385,9 @@ def ai_workflow_chat(
                 repo_name = new_repo_name
                 app_desc = new_app_desc
                 # Yield updated history and state variables (pass UI outputs through)
-                # NOTE: gemini_key and gemini_model are NOT yielded back
                 yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                       attempts, app_desc, repo_name, generated_code, use_grounding)
+                       attempts, app_desc, repo_name, generated_code, use_grounding,
+                       gemini_api_key, gemini_model)
                 # No return needed
 
             elif create_match:
@@ -393,9 +397,9 @@ def ai_workflow_chat(
                  state = STATE_CREATING_SPACE # Transition state to creation
                  repo_name = new_repo_name # Store the validated repo name
                  # Yield updated history and state variables (pass UI outputs through)
-                 # NOTE: gemini_key and gemini_model are NOT yielded back
                  yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                        attempts, app_desc, repo_name, generated_code, use_grounding)
+                        attempts, app_desc, repo_name, generated_code, use_grounding,
+                        gemini_api_key, gemini_model)
                  # No return needed
 
             elif "create" in message.lower() and not repo_id:
@@ -403,18 +407,18 @@ def ai_workflow_chat(
                 history = add_bot_message(history, "Okay, what should the Space be called? (e.g., `my-awesome-app`)")
                 state = STATE_AWAITING_REPO_NAME # Transition to the state where we wait for the name
                 # Yield updated history and state (pass UI outputs through)
-                # NOTE: gemini_key and gemini_model are NOT yielded back
                 yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                       attempts, app_desc, repo_name, generated_code, use_grounding)
+                       attempts, app_desc, repo_name, generated_code, use_grounding,
+                       gemini_api_key, gemini_model)
                 # No return needed
 
             else:
                 # Command not recognized in IDLE state
                 history = add_bot_message(history, "Command not recognized. Try 'generate me a gradio app called myapp', or 'reset'.")
                 # Yield updated history and current state (pass UI outputs through)
-                # NOTE: gemini_key and gemini_model are NOT yielded back
                 yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                       attempts, app_desc, repo_name, generated_code, use_grounding)
+                       attempts, app_desc, repo_name, generated_code, use_grounding,
+                       gemini_api_key, gemini_model)
                 # No return needed
 
 
@@ -426,9 +430,9 @@ def ai_workflow_chat(
              if not new_repo_name or re.search(r'[^a-zA-Z0-9_-]', new_repo_name) or len(new_repo_name) > 100:
                  history = add_bot_message(history, "Invalid name. Please provide a single word/slug for the Space name (letters, numbers, underscores, hyphens only, max 100 chars).")
                  # Stay in AWAITING_REPO_NAME state and yield message (pass UI outputs through)
-                 # NOTE: gemini_key and gemini_model are NOT yielded back
                  yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                        attempts, app_desc, repo_name, generated_code, use_grounding)
+                        attempts, app_desc, repo_name, generated_code, use_grounding,
+                        gemini_api_key, gemini_model)
                  # No return needed
 
              else:
@@ -437,9 +441,9 @@ def ai_workflow_chat(
                  repo_name = new_repo_name # Store the validated repo name
                  # Yield updated history, state, and repo name. UI outputs remain unchanged for now.
                  # The next click will proceed from the STATE_CREATING_SPACE block.
-                 # NOTE: gemini_key and gemini_model are NOT yielded back
                  yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                        attempts, app_desc, repo_name, generated_code, use_grounding)
+                        attempts, app_desc, repo_name, generated_code, use_grounding,
+                        gemini_api_key, gemini_model)
                  # No return needed
 
         # Note: Each 'elif' block below represents a distinct step in the workflow triggered
@@ -450,9 +454,9 @@ def ai_workflow_chat(
              if not repo_name:
                  history = add_bot_message(history, "Internal error: Repo name missing for creation. Resetting.")
                  # Reset relevant states and UI outputs on critical error
-                 # NOTE: gemini_key and gemini_model are NOT yielded back
                  yield (history, None, STATE_IDLE, "<p>Error creating space.</p>", "", "", 0,
-                       None, None, None, use_grounding) # Pass grounding state through
+                       None, None, None, use_grounding,
+                       gemini_api_key, gemini_model) # Pass grounding state through
                  # No return needed
 
              else:
@@ -464,17 +468,17 @@ def ai_workflow_chat(
                      history = add_bot_message(history, f"✅ Space `{repo_id}` created. Click 'Send' to generate and upload code.")
                      state = STATE_GENERATING_CODE # Transition to the next state
                      # Yield updated state variables and history, and the new iframe HTML
-                     # NOTE: gemini_key and gemini_model are NOT yielded back
                      yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                            attempts, app_desc, repo_name, generated_code, use_grounding) # Pass logs and grounding through
+                            attempts, app_desc, repo_name, generated_code, use_grounding,
+                            gemini_api_key, gemini_model) # Pass logs and grounding through
                      # No return needed
 
                  except Exception as e:
                      history = add_bot_message(history, f"❌ Error creating space: {e}. Click 'reset'.")
                      # Yield error message and reset state on failure
-                     # NOTE: gemini_key and gemini_model are NOT yielded back
                      yield (history, None, STATE_IDLE, "<p>Error creating space.</p>", "", "", 0,
-                           None, None, None, use_grounding) # Pass logs and grounding through
+                           None, None, None, use_grounding,
+                           gemini_api_key, gemini_model) # Pass logs and grounding through
                      # No return needed
 
 
@@ -493,9 +497,9 @@ Return **only** the python code block for `app.py`. Do not include any extra tex
                  if use_grounding:
                      history = add_bot_message(history, "(Using Grounding with Google Search)")
                  # Yield to show message before the potentially time-consuming API call
-                 # NOTE: gemini_key and gemini_model are NOT yielded back
                  yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                        attempts, app_desc, repo_name, generated_code, use_grounding)
+                        attempts, app_desc, repo_name, generated_code, use_grounding,
+                        gemini_api_key, gemini_model)
 
                  # Perform the Gemini API call to generate code, optionally using grounding
                  # Use the gemini_api_key and gemini_model inputs directly (which come from state)
@@ -516,17 +520,17 @@ Return **only** the python code block for `app.py`. Do not include any extra tex
                  state = STATE_UPLOADING_APP_PY # Transition to the upload state
                  generated_code = code # Store the generated code in the state variable for the next step
                  # Yield updated state variables and history (pass UI outputs and other states through)
-                 # NOTE: gemini_key and gemini_model are NOT yielded back
                  yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                        attempts, app_desc, repo_name, generated_code, use_grounding)
+                        attempts, app_desc, repo_name, generated_code, use_grounding,
+                        gemini_api_key, gemini_model)
                  # No return needed
 
              except Exception as e:
                  history = add_bot_message(history, f"❌ Error generating code: {e}. Click 'reset'.")
                  # Yield error message and reset state on failure
-                 # NOTE: gemini_key and gemini_model are NOT yielded back
                  yield (history, None, STATE_IDLE, updated_preview, updated_run, updated_build, 0,
-                       None, None, None, use_grounding)
+                       None, None, None, use_grounding,
+                       gemini_api_key, gemini_model)
                  # No return needed
 
 
@@ -535,17 +539,17 @@ Return **only** the python code block for `app.py`. Do not include any extra tex
              code_to_upload = generated_code
              if not code_to_upload:
                   history = add_bot_message(history, "Internal error: No code to upload. Resetting.")
-                  # NOTE: gemini_key and gemini_model are NOT yielded back
                   yield (history, None, STATE_IDLE, updated_preview, updated_run, updated_build, 0,
-                        None, None, None, use_grounding)
+                        None, None, None, use_grounding,
+                        gemini_api_key, gemini_model)
                   # No return needed
 
              else:
                  history = add_bot_message(history, "☁️ Uploading `app.py`...")
                  # Yield to show message before the upload action (pass UI outputs and states through)
-                 # NOTE: gemini_key and gemini_model are NOT yielded back
                  yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                        attempts, app_desc, repo_name, generated_code, use_grounding)
+                        attempts, app_desc, repo_name, generated_code, use_grounding,
+                        gemini_api_key, gemini_model)
 
                  try:
                      # Perform the file upload action
@@ -554,26 +558,26 @@ Return **only** the python code block for `app.py`. Do not include any extra tex
                      state = STATE_GENERATING_REQUIREMENTS # Transition state
                      generated_code = None # Clear the stored code after use to free memory/state space
                      # Yield updated state variables and history (pass UI outputs and other states through)
-                     # NOTE: gemini_key and gemini_model are NOT yielded back
                      yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                            attempts, app_desc, repo_name, generated_code, use_grounding)
+                            attempts, app_desc, repo_name, generated_code, use_grounding,
+                            gemini_api_key, gemini_model)
                      # No return needed
 
                  except Exception as e:
                      history = add_bot_message(history, f"❌ Error uploading `app.py`: {e}. Click 'reset'.")
                      # Yield error message and reset state on failure
-                     # NOTE: gemini_key and gemini_model are NOT yielded back
                      yield (history, None, STATE_IDLE, updated_preview, updated_run, updated_build, 0,
-                           None, None, None, use_grounding)
+                           None, None, None, use_grounding,
+                           gemini_api_key, gemini_model)
                      # No return needed
 
 
         elif state == STATE_GENERATING_REQUIREMENTS:
              history = add_bot_message(history, "📄 Generating `requirements.txt`...")
              # Yield to show message before generating requirements (pass UI outputs and states through)
-             # NOTE: gemini_key and gemini_model are NOT yielded back
              yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                    attempts, app_desc, repo_name, generated_code, use_grounding)
+                    attempts, app_desc, repo_name, generated_code, use_grounding,
+                    gemini_api_key, gemini_model)
 
              # Logic to determine required packages based on SDK and keywords in the app description
              reqs_list = ["gradio"] if space_sdk == "gradio" else ["streamlit"]
@@ -614,9 +618,9 @@ Return **only** the python code block for `app.py`. Do not include any extra tex
              state = STATE_UPLOADING_REQUIREMENTS # Transition state
              generated_code = reqs_content # Store requirements content
              # Yield updated state variables and history (pass UI outputs and other states through)
-             # NOTE: gemini_key and gemini_model are NOT yielded back
              yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                    attempts, app_desc, repo_name, generated_code, use_grounding)
+                    attempts, app_desc, repo_name, generated_code, use_grounding,
+                    gemini_api_key, gemini_model)
                      # No return needed
 
 
@@ -625,17 +629,17 @@ Return **only** the python code block for `app.py`. Do not include any extra tex
             reqs_content_to_upload = generated_code
             if not reqs_content_to_upload:
                  history = add_bot_message(history, "Internal error: No requirements content to upload. Resetting.")
-                 # NOTE: gemini_key and gemini_model are NOT yielded back
                  yield (history, None, STATE_IDLE, updated_preview, updated_run, updated_build, 0,
-                        None, None, None, use_grounding)
+                        None, None, None, use_grounding,
+                        gemini_api_key, gemini_model)
                  # No return needed
 
             else:
                  history = add_bot_message(history, "☁️ Uploading `requirements.txt`...")
                  # Yield message before upload (pass UI outputs and states through)
-                 # NOTE: gemini_key and gemini_model are NOT yielded back
                  yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                        attempts, app_desc, repo_name, generated_code, use_grounding)
+                        attempts, app_desc, repo_name, generated_code, use_grounding,
+                        gemini_api_key, gemini_model)
 
                  try:
                      # Perform requirements file upload
@@ -644,25 +648,25 @@ Return **only** the python code block for `app.py`. Do not include any extra tex
                      state = STATE_GENERATING_README # Transition state
                      generated_code = None # Clear content after use
                      # Yield updated state variables and history (pass UI outputs and other states through)
-                     # NOTE: gemini_key and gemini_model are NOT yielded back
                      yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                            attempts, app_desc, repo_name, generated_code, use_grounding)
+                            attempts, app_desc, repo_name, generated_code, use_grounding,
+                            gemini_api_key, gemini_model)
                      # No return needed
 
                  except Exception as e:
                      history = add_bot_message(history, f"❌ Error uploading `requirements.txt`: {e}. Click 'reset'.")
                      # Yield error message and reset state on failure
-                     # NOTE: gemini_key and gemini_model are NOT yielded back
                      yield (history, None, STATE_IDLE, updated_preview, updated_run, updated_build, 0,
-                           None, None, None, use_grounding)
+                           None, None, None, use_grounding,
+                           gemini_api_key, gemini_model)
                      # No return needed
 
         elif state == STATE_GENERATING_README:
             history = add_bot_message(history, "📝 Generating `README.md`...")
             # Yield message before generating README (pass UI outputs and states through)
-            # NOTE: gemini_key and gemini_model are NOT yielded back
             yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                   attempts, app_desc, repo_name, generated_code, use_grounding)
+                   attempts, app_desc, repo_name, generated_code, use_grounding,
+                   gemini_api_key, gemini_model)
 
             # Generate simple README content with Space metadata header
             readme_title = repo_name if repo_name else "My Awesome Space"
@@ -689,9 +693,9 @@ This Space was automatically generated by an AI workflow using Google Gemini and
             state = STATE_UPLOADING_README # Transition state
             generated_code = readme_content # Store README content
             # Yield updated state variables and history (pass UI outputs and other states through)
-            # NOTE: gemini_key and gemini_model are NOT yielded back
             yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                   attempts, app_desc, repo_name, generated_code, use_grounding)
+                   attempts, app_desc, repo_name, generated_code, use_grounding,
+                   gemini_api_key, gemini_model)
             # No return needed
 
 
@@ -700,17 +704,17 @@ This Space was automatically generated by an AI workflow using Google Gemini and
             readme_content_to_upload = generated_code
             if not readme_content_to_upload:
                  history = add_bot_message(history, "Internal error: No README content to upload. Resetting.")
-                 # NOTE: gemini_key and gemini_model are NOT yielded back
                  yield (history, None, STATE_IDLE, updated_preview, updated_run, updated_build, 0,
-                        None, None, None, use_grounding)
+                        None, None, None, use_grounding,
+                        gemini_api_key, gemini_model)
                  # No return needed
 
             else:
                  history = add_bot_message(history, "☁️ Uploading `README.md`...")
                  # Yield message before upload (pass UI outputs and states through)
-                 # NOTE: gemini_key and gemini_model are NOT yielded back
                  yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                        attempts, app_desc, repo_name, generated_code, use_grounding)
+                        attempts, app_desc, repo_name, generated_code, use_grounding,
+                        gemini_api_key, gemini_model)
 
                  try:
                      # Perform README file upload
@@ -719,25 +723,25 @@ This Space was automatically generated by an AI workflow using Google Gemini and
                      state = STATE_CHECKING_LOGS_BUILD # Transition to checking build logs
                      generated_code = None # Clear content after use
                      # Yield updated state variables and history (pass UI outputs and other states through)
-                     # NOTE: gemini_key and gemini_model are NOT yielded back
                      yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                            attempts, app_desc, repo_name, generated_code, use_grounding)
+                            attempts, app_desc, repo_name, generated_code, use_grounding,
+                            gemini_api_key, gemini_model)
                      # No return needed
 
                  except Exception as e:
                      history = add_bot_message(history, f"❌ Error uploading `README.md`: {e}. Click 'reset'.")
                      # Yield error message and reset state on failure
-                     # NOTE: gemini_key and gemini_model are NOT yielded back
                      yield (history, None, STATE_IDLE, updated_preview, updated_run, updated_build, 0,
-                           None, None, None, use_grounding)
+                           None, None, None, use_grounding,
+                           gemini_api_key, gemini_model)
                      # No return needed
 
         elif state == STATE_CHECKING_LOGS_BUILD:
              history = add_bot_message(history, "🔍 Fetching build logs...")
              # Yield message before fetching logs (which includes a delay) (pass UI outputs and states through)
-             # NOTE: gemini_key and gemini_model are NOT yielded back
              yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                    attempts, app_desc, repo_name, generated_code, use_grounding)
+                    attempts, app_desc, repo_name, generated_code, use_grounding,
+                    gemini_api_key, gemini_model)
 
              # Fetch build logs from HF Space
              build_logs_text = get_build_logs_action(repo_id, hf_profile, hf_token)
@@ -748,27 +752,27 @@ This Space was automatically generated by an AI workflow using Google Gemini and
                   history = add_bot_message(history, "⚠️ Build logs indicate potential issues. Please inspect above. Click 'Send' to check container logs (app might still start despite build warnings).")
                   state = STATE_CHECKING_LOGS_RUN # Transition even on build error, to see if container starts
                   # Yield updated state, logs, and variables
-                  # NOTE: gemini_key and gemini_model are NOT yielded back
                   yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                         attempts, app_desc, repo_name, generated_code, use_grounding)
+                         attempts, app_desc, repo_name, generated_code, use_grounding,
+                         gemini_api_key, gemini_model)
                   # No return needed
 
              else:
                   history = add_bot_message(history, "✅ Build logs fetched. Click 'Send' to check container logs.")
                   state = STATE_CHECKING_LOGS_RUN # Transition to next log check
                   # Yield updated state, logs, and variables
-                  # NOTE: gemini_key and gemini_model are NOT yielded back
                   yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                         attempts, app_desc, repo_name, generated_code, use_grounding)
+                         attempts, app_desc, repo_name, generated_code, use_grounding,
+                         gemini_api_key, gemini_model)
                   # No return needed
 
 
         elif state == STATE_CHECKING_LOGS_RUN:
              history = add_bot_message(history, "🔍 Fetching container logs...")
              # Yield message before fetching logs (includes a delay) (pass UI outputs and states through)
-             # NOTE: gemini_key and gemini_model are NOT yielded back
              yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                    attempts, app_desc, repo_name, generated_code, use_grounding)
+                    attempts, app_desc, repo_name, generated_code, use_grounding,
+                    gemini_api_key, gemini_model)
 
              # Fetch container logs from HF Space
              container_logs_text = get_container_logs_action(repo_id, hf_profile, hf_token)
@@ -780,9 +784,9 @@ This Space was automatically generated by an AI workflow using Google Gemini and
                   history = add_bot_message(history, f"❌ Errors detected in container logs. Attempting debug fix #{attempts}/{MAX_DEBUG_ATTEMPTS}. Click 'Send' to proceed.")
                   state = STATE_DEBUGGING_CODE # Transition to the debugging state
                   # Yield updated state, logs, attempts, and variables
-                  # NOTE: gemini_key and gemini_model are NOT yielded back
                   yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                         attempts, app_desc, repo_name, generated_code, use_grounding)
+                         attempts, app_desc, repo_name, generated_code, use_grounding,
+                         gemini_api_key, gemini_model)
                   # No return needed
 
              elif ("error" in updated_run.lower() or "exception" in updated_run.lower()) and attempts >= MAX_DEBUG_ATTEMPTS:
@@ -790,9 +794,9 @@ This Space was automatically generated by an AI workflow using Google Gemini and
                   history = add_bot_message(history, f"❌ Errors detected in container logs. Max debug attempts ({MAX_DEBUG_ATTEMPTS}) reached. Please inspect logs manually or click 'reset'.")
                   state = STATE_COMPLETE # Workflow ends on failure after attempts
                   # Yield updated state, logs, attempts, and variables
-                  # NOTE: gemini_key and gemini_model are NOT yielded back
                   yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                         attempts, app_desc, repo_name, generated_code, use_grounding)
+                         attempts, app_desc, repo_name, generated_code, use_grounding,
+                         gemini_api_key, gemini_model)
                   # No return needed
 
              else:
@@ -800,9 +804,9 @@ This Space was automatically generated by an AI workflow using Google Gemini and
                   history = add_bot_message(history, "✅ App appears to be running successfully! Check the iframe above. Click 'reset' to start a new project.")
                   state = STATE_COMPLETE # Workflow ends on success
                   # Yield updated state, logs, attempts, and variables
-                  # NOTE: gemini_key and gemini_model are NOT yielded back
                   yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                         attempts, app_desc, repo_name, generated_code, use_grounding)
+                         attempts, app_desc, repo_name, generated_code, use_grounding,
+                         gemini_api_key, gemini_model)
                   # No return needed
 
 
@@ -811,9 +815,9 @@ This Space was automatically generated by an AI workflow using Google Gemini and
              if use_grounding:
                   history = add_bot_message(history, "(Using Grounding with Google Search)")
              # Yield message before Gemini API call (pass UI outputs and states through)
-             # NOTE: gemini_key and gemini_model are NOT yielded back
              yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                    attempts, app_desc, repo_name, generated_code, use_grounding)
+                    attempts, app_desc, repo_name, generated_code, use_grounding,
+                    gemini_api_key, gemini_model)
 
              # Construct prompt for Gemini including the container logs
              debug_prompt = f"""
@@ -848,17 +852,17 @@ Return **only** the python code block for app.py. Do not include any extra text,
                   state = STATE_UPLOADING_FIXED_APP_PY # Transition to the upload state for the fix
                   generated_code = fix_code # Store the generated fix code
                   # Yield updated state, code, and variables (pass UI outputs and states through)
-                  # NOTE: gemini_key and gemini_model are NOT yielded back
                   yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                         attempts, app_desc, repo_name, generated_code, use_grounding)
+                         attempts, app_desc, repo_name, generated_code, use_grounding,
+                         gemini_api_key, gemini_model)
                   # No return needed
 
              except Exception as e:
                   history = add_bot_message(history, f"❌ Error generating debug code: {e}. Click 'reset'.")
                   # Yield error message and reset state on failure
-                  # NOTE: gemini_key and gemini_model are NOT yielded back
                   yield (history, None, STATE_IDLE, updated_preview, updated_run, updated_build, 0,
-                         None, None, None, use_grounding)
+                         None, None, None, use_grounding,
+                         gemini_api_key, gemini_model)
                   # No return needed
 
         elif state == STATE_UPLOADING_FIXED_APP_PY:
@@ -866,17 +870,17 @@ Return **only** the python code block for app.py. Do not include any extra text,
              fixed_code_to_upload = generated_code
              if not fixed_code_to_upload:
                   history = add_bot_message(history, "Internal error: No fixed code available to upload. Resetting.")
-                  # NOTE: gemini_key and gemini_model are NOT yielded back
                   yield (history, None, STATE_IDLE, updated_preview, updated_run, updated_build, 0,
-                        None, None, None, use_grounding)
+                        None, None, None, use_grounding,
+                        gemini_api_key, gemini_model)
                   # No return needed
 
              else:
                   history = add_bot_message(history, "☁️ Uploading fixed `app.py`...")
                   # Yield message before upload (pass UI outputs and states through)
-                  # NOTE: gemini_key and gemini_model are NOT yielded back
                   yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                         attempts, app_desc, repo_name, generated_code, use_grounding)
+                         attempts, app_desc, repo_name, generated_code, use_grounding,
+                         gemini_api_key, gemini_model)
 
                   try:
                       # Perform the upload of the fixed app.py
@@ -885,25 +889,25 @@ Return **only** the python code block for app.py. Do not include any extra text,
                       state = STATE_CHECKING_LOGS_RUN # Go back to checking run logs after uploading the fix
                       generated_code = None # Clear code after use
                       # Yield updated state, code, and variables (pass UI outputs and states through)
-                      # NOTE: gemini_key and gemini_model are NOT yielded back
                       yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                             attempts, app_desc, repo_name, generated_code, use_grounding)
+                             attempts, app_desc, repo_name, generated_code, use_grounding,
+                             gemini_api_key, gemini_model)
                       # No return needed
 
                   except Exception as e:
                       history = add_bot_message(history, f"❌ Error uploading fixed `app.py`: {e}. Click 'reset'.")
                       # Yield error message and reset state on failure
-                      # NOTE: gemini_key and gemini_model are NOT yielded back
                       yield (history, None, STATE_IDLE, updated_preview, updated_run, updated_build, 0,
-                            None, None, None, use_grounding)
+                            None, None, None, use_grounding,
+                            gemini_api_key, gemini_model)
                       # No return needed
 
         elif state == STATE_COMPLETE:
              # If in the complete state, the workflow is finished for this project.
              # Subsequent clicks just add user messages; we simply yield the current state.
-             # NOTE: gemini_key and gemini_model are NOT yielded back
              yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                    attempts, app_desc, repo_name, generated_code, use_grounding)
+                    attempts, app_desc, repo_name, generated_code, use_grounding,
+                    gemini_api_key, gemini_model)
              # No return needed
 
 
@@ -913,9 +917,9 @@ Return **only** the python code block for app.py. Do not include any extra text,
         history = add_bot_message(history, error_message)
         print(f"Critical Error in state {state}: {e}") # Log the error for debugging purposes
         # Yield an error state and reset essential workflow variables on critical failure
-        # NOTE: gemini_key and gemini_model are NOT yielded back
         yield (history, None, STATE_IDLE, updated_preview, updated_run, updated_build, 0,
-               None, None, None, use_grounding) # Include use_grounding
+               None, None, None, use_grounding,
+               gemini_api_key, gemini_model) # Include use_grounding and Gemini states
         # No return needed after yield
 
 
@@ -1090,8 +1094,9 @@ with gr.Blocks(title="AI-Powered HF Space App Builder") as ai_builder_tab:
                     repo_id, workflow, # Update workflow State variables
                     iframe, run_txt, build_txt, # Update UI component outputs
                     debug_attempts, app_description, repo_name_state, generated_code_state, # Update other State variables
-                    use_grounding_state # Update the grounding state output
-                    # REMOVED: gemini_key, gemini_model from outputs
+                    use_grounding_state, # Update the grounding state output
+                    # ADDED BACK: gemini_key, gemini_model to outputs to ensure state consistency
+                    gemini_key, gemini_model
                 ]
             ).success( # Chain a .success() event to run *after* the .click() handler completes without error
                  # Clear the user input textbox after the message is sent and processed