diff --git "a/app.py" "b/app.py"
--- "a/app.py"
+++ "b/app.py"
@@ -4,6 +4,8 @@ import time
 import json
 import io
 import requests
+import logging
+from typing import List, Dict, Any, Tuple, Optional, Literal, Generator
 
 import gradio as gr
 import google.generativeai as genai
@@ -11,11 +13,26 @@ from google.generativeai import types # Import types for configuration and tools
 
 from huggingface_hub import create_repo, list_models, upload_file, constants
 from huggingface_hub.utils import build_hf_headers, get_session, hf_raise_for_status
+from requests.adapters import HTTPAdapter
+from urllib3.util.retry import Retry
+
+
+# --- Configure Logging ---
+# Replace print() statements with logging
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+# You could add a file handler here for persistent logs if needed, but console is fine for Spaces
+
+
+# --- Configure Hugging Face API Retries ---
+# Added retry strategy to make HF API calls more robust to transient errors
+retry_strategy = Retry(total=5, backoff_factor=1, status_forcelist=[429, 500, 502, 503, 504]) # Define retry strategy for specific HTTP codes
+adapter = HTTPAdapter(max_retries=retry_strategy)
+session = get_session() # Get the session object used internally by huggingface_hub
+session.mount("http://", adapter)
+session.mount("https://", adapter)
 
-# Removed the debugging print that attempts to read GOOGLE_API_KEY from environment
 
 # --- Define Gemini Model Information ---
-# Dictionary mapping internal model name to (Display Name, Description)
 GEMINI_MODELS = {
     "gemini-1.5-flash": ("Gemini 1.5 Flash", "Fast and versatile performance across a diverse variety of tasks."),
     "gemini-1.5-pro": ("Gemini 1.5 Pro", "Complex reasoning tasks requiring more intelligence."),
@@ -28,10 +45,9 @@ GEMINI_MODELS = {
 }
 
 # Create the list of choices for the Gradio Radio component
-# Format is (Display Name, Internal Name)
 GEMINI_MODEL_CHOICES = [(display_name, internal_name) for internal_name, (display_name, description) in GEMINI_MODELS.items()]
-# Define the default model to be selected
-DEFAULT_GEMINI_MODEL = "gemini-1.5-flash" # Ensure this key exists in GEMINI_MODELS
+DEFAULT_GEMINI_MODEL = "gemini-1.5-flash"
+
 
 # --- Helper functions for Hugging Face integration ---
 
@@ -41,6 +57,7 @@ def show_profile(profile: gr.OAuthProfile | None) -> str:
         return "*Not logged in.*"
     return f"✅ Logged in as **{profile.username}**"
 
+# list_private_models function is not used in the main workflow, kept as is.
 def list_private_models(
     profile: gr.OAuthProfile | None,
     oauth_token: gr.OAuthToken | None
@@ -55,15 +72,17 @@ def list_private_models(
         ]
         return "No models found." if not models else "Models:\n\n" + "\n - ".join(models)
     except Exception as e:
-        # Catching generic exception is acceptable for helper functions
+        logging.error(f"Error listing models: {e}")
         return f"Error listing models: {e}"
 
-def create_space_action(repo_name: str, sdk: str, profile: gr.OAuthProfile, token: gr.OAuthToken):
+def create_space_action(repo_name: str, sdk: str, profile: gr.OAuthProfile, token: gr.OAuthToken) -> Tuple[str, str]:
     """Creates a new Hugging Face Space repository."""
     if not profile or not token:
+        # This should ideally not happen if button logic is correct, but kept as safeguard
         raise ValueError("Hugging Face profile or token is missing.")
     repo_id = f"{profile.username}/{repo_name}"
     try:
+        logging.info(f"Attempting to create Space: {repo_id} with SDK: {sdk}")
         create_repo(
             repo_id=repo_id,
             token=token.token,
@@ -73,9 +92,14 @@ def create_space_action(repo_name: str, sdk: str, profile: gr.OAuthProfile, toke
         )
         url    = f"https://huggingface.co/spaces/{repo_id}"
         iframe = f'<iframe src="{url}" width="100%" height="500px"></iframe>'
+        logging.info(f"Successfully created/verified Space: {repo_id}")
         return repo_id, iframe
     except Exception as e:
-        raise RuntimeError(f"Failed to create Space `{repo_id}`: {e}")
+        logging.error(f"Failed to create Space {repo_id}: {e}")
+        # Catch specific HTTP errors from huggingface_hub if possible
+        if isinstance(e, requests.exceptions.HTTPError):
+            raise RuntimeError(f"HF API Error creating Space `{repo_id}`: {e.response.status_code} {e.response.reason}") from e
+        raise RuntimeError(f"Failed to create Space `{repo_id}`: {e}") from e # Re-raise as RuntimeError
 
 def upload_file_to_space_action(
     file_obj: io.StringIO, # Specify type hint for clarity
@@ -84,10 +108,11 @@ def upload_file_to_space_action(
     profile: gr.OAuthProfile,
     token: gr.OAuthToken
 ) -> None:
-    """Uploads a file to a Huging Face Space repository."""
+    """Uploads a file to a Hugging Face Space repository."""
     if not (profile and token and repo_id):
         raise ValueError("Hugging Face profile, token, or repo_id is missing.")
     try:
+        logging.info(f"Attempting to upload file: {path_in_repo} to Space: {repo_id}")
         upload_file(
             path_or_fileobj=file_obj,
             path_in_repo=path_in_repo,
@@ -95,16 +120,23 @@ def upload_file_to_space_action(
             token=token.token,
             repo_type="space"
         )
+        logging.info(f"Successfully uploaded file: {path_in_repo} to Space: {repo_id}")
     except Exception as e:
-        raise RuntimeError(f"Failed to upload `{path_in_repo}` to `{repo_id}`: {e}")
+        logging.error(f"Failed to upload {path_in_repo} to {repo_id}: {e}")
+        if isinstance(e, requests.exceptions.HTTPError):
+             raise RuntimeError(f"HF API Error uploading {path_in_repo} to `{repo_id}`: {e.response.status_code} {e.response.reason}") from e
+        raise RuntimeError(f"Failed to upload `{path_in_repo}` to `{repo_id}`: {e}") from e
+
 
 def _fetch_space_logs_level(repo_id: str, level: str, token: str) -> str:
     """Fetches build or run logs for a Space."""
     if not repo_id or not token:
-         return f"Cannot fetch {level} logs: repo_id or token missing."
+         logging.warning(f"Cannot fetch {level} logs: repo_id or token missing.")
+         return f"Cannot fetch {level} logs: log in and create a Space first."
     jwt_url  = f"{constants.ENDPOINT}/api/spaces/{repo_id}/jwt"
     try:
-        r = get_session().get(jwt_url, headers=build_hf_headers(token=token))
+        logging.info(f"Attempting to fetch {level} logs for Space: {repo_id}")
+        r = get_session().get(jwt_url, headers=build_hf_headers(token=token), timeout=10) # Added timeout
         hf_raise_for_status(r) # Raise HTTPError for bad responses (4xx or 5xx)
         jwt = r.json()["token"]
         logs_url = f"https://api.hf.space/v1/{repo_id}/logs/{level}"
@@ -128,10 +160,17 @@ def _fetch_space_logs_level(repo_id: str, level: str, token: str) -> str:
                         count += 1
                 except json.JSONDecodeError:
                     # Skip lines that aren't valid JSON events
+                    logging.warning(f"Skipping non-JSON log line for {repo_id} ({level}): {payload.decode()}")
                     continue
-        return "\n".join(lines) if lines else f"No {level} logs found."
+        log_output = "\n".join(lines) if lines else f"No {level} logs found."
+        logging.info(f"Successfully fetched {count} {level} log lines for {repo_id}")
+        return log_output
     except Exception as e:
-        # Catching generic exception is acceptable for helper functions
+        logging.error(f"Error fetching {level} logs for {repo_id}: {e}")
+        if isinstance(e, requests.exceptions.HTTPError):
+             return f"Error fetching {level} logs for `{repo_id}`: {e.response.status_code} {e.response.reason}"
+        if isinstance(e, requests.exceptions.Timeout):
+             return f"Timeout fetching {level} logs for `{repo_id}`. Space might be starting slowly."
         return f"Error fetching {level} logs for `{repo_id}`: {e}"
 
 
@@ -157,141 +196,1228 @@ def get_container_logs_action(repo_id, profile, token):
 def configure_gemini(api_key: str | None, model_name: str | None) -> str:
     """Configures the Gemini API and checks if the model is accessible."""
     # Check for empty string "" as well as None
-    if not api_key:
+    if not isinstance(api_key, str) or not api_key.strip():
+        logging.info("Gemini API key is not set.")
         return "⚠️ Gemini API key is not set."
     # Check if model_name is None or not a valid key in GEMINI_MODELS
     if not model_name or model_name not in GEMINI_MODELS:
+        logging.warning(f"Invalid Gemini model selected: {model_name}")
         return "⚠️ Please select a valid Gemini model."
     try:
+        logging.info(f"Attempting to configure Gemini with model: {model_name}")
         genai.configure(api_key=api_key)
         # Attempt a simple call to verify credentials and model availability
         # This will raise an exception if the key is invalid or model not found
         genai.GenerativeModel(model_name).generate_content("ping", stream=False)
         # This message indicates the API call *for configuration check* was successful
+        logging.info(f"Gemini configured successfully with model: {model_name}")
         return f"✅ Gemini configured successfully with **{GEMINI_MODELS[model_name][0]}**."
     except Exception as e:
         # This message indicates the API call *for configuration check* failed
+        logging.error(f"Error configuring Gemini with model {model_name}: {e}")
+        # Catch specific Gemini errors if possible (e.g., authentication errors)
         return f"❌ Error configuring Gemini: {e}"
 
 def get_model_description(model_name: str | None) -> str:
     """Retrieves the description for a given model name."""
     if model_name is None or model_name not in GEMINI_MODELS:
         return "Select a model to see its description."
-    # Use .get with a default value to handle cases where the key might not be found
     return GEMINI_MODELS.get(model_name, (model_name, "No description available."))[1]
 
 
 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
-    # 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}")
+    # These checks are crucial - they will raise an error *before* the API call if prereqs aren't met
+    if not isinstance(api_key, str) or not api_key.strip():
+        raise ValueError("Gemini API key is empty or invalid.")
+    if not model_name or model_name not in GEMINI_MODELS:
+         raise ValueError(f"Gemini model '{model_name}' is invalid or not selected.")
+
     try:
-        genai.configure(api_key=api_key)
+        logging.info(f"Calling Gemini model '{model_name}' (Grounding: {use_grounding}) with prompt (first 50 chars): '{prompt[:50]}...'")
+        genai.configure(api_key=api_key) # Re-configure just in case
         model = genai.GenerativeModel(model_name)
 
-        # Define tools for grounding if requested.
-        # Using genai.types.GoogleSearch() is recommended for Gemini 2.0+
-        # and is backwards compatible with 1.5 for retrieval.
         tools_config = [types.Tool(google_search=types.GoogleSearch())] if use_grounding else None
 
-        # Using generate_content and stream=False for simplicity here
         response = model.generate_content(
             prompt,
-            stream=False,
-            tools=tools_config # Pass the tools configuration
+            stream=False, # Using stream=False for simplicity in this workflow
+            tools=tools_config,
+            request_options={'timeout': 120} # Added timeout for API call
         )
-        # Check if response is blocked
+
         if response.prompt_feedback and response.prompt_feedback.block_reason:
+             logging.warning(f"Gemini API call blocked: {response.prompt_feedback.block_reason}")
              raise RuntimeError(f"Gemini API call blocked: {response.prompt_feedback.block_reason}")
+
         if not response.candidates:
-             # Check for safety ratings if no candidates are returned but not blocked
              if response.prompt_feedback and response.prompt_feedback.safety_ratings:
                   ratings = "; ".join([f"{r.category}: {r.probability}" for r in response.prompt_feedback.safety_ratings])
+                  logging.warning(f"Gemini API call returned no candidates. Safety ratings: {ratings}")
                   raise RuntimeError(f"Gemini API call returned no candidates. Safety ratings: {ratings}")
              else:
+                  logging.warning("Gemini API call returned no candidates.")
                   raise RuntimeError("Gemini API call returned no candidates.")
 
-
-        # If response.candidates is not empty, get the text
-        # Using response.text is a convenient way to get text from the first candidate part
-        return response.text or "" # Return empty string if no text
+        generated_text = response.text or ""
+        logging.info(f"Gemini API call successful. Generated text length: {len(generated_text)}")
+        return generated_text
 
     except Exception as e:
+        logging.error(f"Gemini API call failed: {e}")
         # Re-raising as RuntimeError for the workflow to catch and manage
-        raise RuntimeError(f"Gemini API call failed: {e}")
+        raise RuntimeError(f"Gemini API call failed: {e}") from e
 
 
 # --- AI workflow logic (State Machine) ---
 
-# Define States for the workflow
-STATE_IDLE = "idle"
-STATE_AWAITING_REPO_NAME = "awaiting_repo_name"
-STATE_CREATING_SPACE = "creating_space"
-STATE_GENERATING_CODE = "generating_code"
-STATE_UPLOADING_APP_PY = "uploading_app_py"
-STATE_GENERATING_REQUIREMENTS = "generating_requirements"
-STATE_UPLOADING_REQUIREMENTS = "uploading_requirements"
-STATE_GENERATING_README = "generating_readme"
-STATE_UPLOADING_README = "uploading_readme"
-STATE_CHECKING_LOGS_BUILD = "checking_logs_build"
-STATE_CHECKING_LOGS_RUN = "checking_logs_run"
-STATE_DEBUGGING_CODE = "debugging_code"
-STATE_UPLOADING_FIXED_APP_PY = "uploading_fixed_app_py"
-STATE_COMPLETE = "complete"
+# Define States for the workflow using Literal for type safety
+WorkflowState = Literal[
+    "idle", "awaiting_repo_name", "creating_space", "generating_code",
+    "uploading_app_py", "generating_requirements", "uploading_requirements",
+    "generating_readme", "uploading_readme", "checking_logs_build",
+    "checking_logs_run", "debugging_code", "uploading_fixed_app_py", "complete"
+]
+
+STATE_IDLE: WorkflowState = "idle"
+STATE_AWAITING_REPO_NAME: WorkflowState = "awaiting_repo_name"
+STATE_CREATING_SPACE: WorkflowState = "creating_space"
+STATE_GENERATING_CODE: WorkflowState = "generating_code"
+STATE_UPLOADING_APP_PY: WorkflowState = "uploading_app_py"
+STATE_GENERATING_REQUIREMENTS: WorkflowState = "generating_requirements"
+STATE_UPLOADING_REQUIREMENTS: WorkflowState = "uploading_requirements"
+STATE_GENERATING_README: WorkflowState = "generating_readme"
+STATE_UPLOADING_README: WorkflowState = "uploading_readme"
+STATE_CHECKING_LOGS_BUILD: WorkflowState = "checking_logs_build"
+STATE_CHECKING_LOGS_RUN: WorkflowState = "checking_logs_run"
+STATE_DEBUGGING_CODE: WorkflowState = "debugging_code"
+STATE_UPLOADING_FIXED_APP_PY: WorkflowState = "uploading_fixed_app_py"
+STATE_COMPLETE: WorkflowState = "complete"
 
 MAX_DEBUG_ATTEMPTS = 3 # Limit the number of automatic debug attempts
 
+# Helper function to add a new assistant message to the chatbot history.
 def add_bot_message(history: list[dict], bot_message: str) -> list[dict]:
-    """Helper to add a new assistant message to the chatbot history."""
-    history.append({"role": "assistant", "content": bot_message})
-    return history
+    # Make a copy to avoid modifying history in place if needed later, though generator pattern usually handles this
+    new_history = list(history)
+    new_history.append({"role": "assistant", "content": bot_message})
+    logging.info(f"Added bot message: {bot_message[:100]}...")
+    return new_history
 
 # Add an initial welcome message to the chatbot (defined outside Blocks to be called by load chain)
-def greet():
-    # Updated welcome message to reflect the change in API key handling
+def greet() -> List[Dict[str, str]]:
+    logging.info("Generating initial welcome message.")
     return [{"role": "assistant", "content": "Welcome! Please log in to Hugging Face and provide your Google AI Studio API key to start building Spaces. Once ready, type 'generate me a gradio app called myapp' or 'create' to begin."}]
 
+
 # Helper function to update send button interactivity based on prereqs
-# This function has the clean signature it expects. Lambdas handle Gradio's argument passing.
-def check_send_button_ready(hf_profile: gr.OAuthProfile | None, hf_token: gr.OAuthToken | None, gemini_key: str | None, gemini_model: str | None) -> gr.update:
+# This function has the clean signature it expects.
+def check_send_button_ready(
+    hf_profile: gr.OAuthProfile | None,
+    hf_token: gr.OAuthToken | None,
+    gemini_key: str | None,
+    gemini_model: str | None,
+    workflow_state: WorkflowState # Also depend on workflow state
+) -> gr.Button.update: # Use specific component update
     """Checks if HF login and Gemini configuration are complete and returns update for button interactivity."""
-    # --- START ENHANCED DEBUGGING LOGS ---
-    print("\n--- check_send_button_ready START ---")
-    print(f"  Received hf_profile: {hf_profile is not None}")
-    print(f"  Received hf_token: {hf_token is not None}")
-    # For api_key, print part of the key if not None/empty for verification
-    api_key_display = gemini_key[:5] if isinstance(gemini_key, str) and gemini_key else ('Empty String' if isinstance(gemini_key, str) and gemini_key == "" else 'None')
-    print(f"  Received gemini_key: Value starts with '{api_key_display}'")
-    print(f"  Received gemini_model: {gemini_model}")
-    # --- END ENHANCED DEBUGGING LOGS ---
+    # Button should NOT be interactive when workflow is running
+    if workflow_state != STATE_IDLE and workflow_state != STATE_AWAITING_REPO_NAME:
+         logging.debug(f"check_send_button_ready: Workflow state is {workflow_state}, disabling button.")
+         return gr.Button.update(interactive=False)
 
     is_logged_in    = hf_profile is not None and hf_token is not None
-    # Use bool() check for simplicity - handles None and "" correctly
-    is_gemini_ready = bool(gemini_key) and bool(gemini_model)
+    # Use strip() to handle cases where key is just whitespace
+    is_gemini_ready = isinstance(gemini_key, str) and bool(gemini_key.strip()) and bool(gemini_model)
 
     is_ready = is_logged_in and is_gemini_ready
-    print(f"check_send_button_ready - HF Ready: {is_logged_in}, Gemini Ready: {is_gemini_ready}, Button Ready: {is_ready}")
-    print("--- check_send_button_ready END ---\n")
+    logging.debug(f"check_send_button_ready - HF Ready: {is_logged_in}, Gemini Ready: {is_gemini_ready}, Button Ready: {is_ready}")
+
+    # Button is interactive only in IDLE or AWAITING_REPO_NAME states AND when prereqs are met
+    return gr.Button.update(interactive=is_ready and (workflow_state == STATE_IDLE or workflow_state == STATE_AWAITING_REPO_NAME))
+
+
+# --- State Handler Functions ---
+# These functions encapsulate the logic for each state.
+# They take all necessary inputs from the main generator's arguments
+# and return the full tuple of outputs required by the generator's yield signature.
+
+WorkflowInputs = Tuple[
+    str, List[Dict[str, str]], Optional[gr.OAuthProfile], Optional[gr.OAuthToken],
+    Optional[str], Optional[str], Optional[str], WorkflowState, str, str, str, str,
+    int, Optional[str], Optional[str], Optional[str], bool
+]
+
+WorkflowOutputs = Tuple[
+    List[Dict[str, str]], Optional[str], WorkflowState, str, str, str,
+    int, Optional[str], Optional[str], Optional[str], bool, Optional[str], Optional[str]
+]
+
+def package_workflow_outputs(
+    history: List[Dict[str, str]],
+    repo_id: Optional[str],
+    state: WorkflowState,
+    updated_preview: str,
+    updated_run: str,
+    updated_build: str,
+    attempts: int,
+    app_desc: Optional[str],
+    repo_name: Optional[str],
+    generated_code: Optional[str],
+    use_grounding: bool,
+    current_gemini_key: Optional[str], # Explicitly include these
+    current_gemini_model: Optional[str] # Explicitly include these
+) -> WorkflowOutputs:
+     """Helper to package all workflow state and UI outputs into the required tuple."""
+     return (history, repo_id, state, updated_preview, updated_run, updated_build,
+             attempts, app_desc, repo_name, generated_code, use_grounding,
+             current_gemini_key, current_gemini_model)
+
+
+def handle_idle(
+    message: str,
+    history: List[Dict[str, str]],
+    hf_profile: gr.OAuthProfile | None,
+    hf_token:   gr.OAuthToken   | None,
+    gemini_api_key: str | None,
+    gemini_model:   str | None,
+    repo_id:  str | None,
+    state: WorkflowState, # Should be STATE_IDLE
+    space_sdk:      str,
+    preview_html:   str,
+    container_logs: str,
+    build_logs:     str,
+    debug_attempts: int,
+    app_description: str | None,
+    repo_name: str | None,
+    generated_code: str | None,
+    use_grounding: bool,
+    *args, # Catch potential extra args
+    **kwargs # Catch potential extra kwargs
+) -> WorkflowOutputs:
+    """Handles logic when in the IDLE state."""
+    logging.info(f"Handling STATE_IDLE with message: {message[:50]}...")
+
+    reset_match = "reset" in message.lower()
+    generate_match = re.search(r'generate (?:me )?(?:a|an) (.+) app called (\w+)', message, re.I)
+    create_match = re.search(r'create (?:a|an)? space called (\w+)', message, re.I) # Simple create command
+
+    current_gemini_key = gemini_api_key # Use the input vars directly
+    current_gemini_model = gemini_model
+
+    if reset_match:
+        logging.info("Reset command received.")
+        history = add_bot_message(history, "Workflow reset.")
+        # Reset relevant states and UI outputs
+        return package_workflow_outputs(
+            history=history, repo_id=None, state=STATE_IDLE,
+            updated_preview="<p>No Space created yet.</p>", updated_run="", updated_build="",
+            attempts=0, app_desc=None, repo_name=None, generated_code=None,
+            use_grounding=use_grounding, current_gemini_key=current_gemini_key, current_gemini_model=current_gemini_model
+        )
+
+    elif generate_match:
+        logging.info("Generate command received.")
+        new_app_desc = generate_match.group(1).strip() # Capture description part
+        new_repo_name = generate_match.group(2).strip() # Capture name part
+        # Perform basic validation on repo name format
+        if not new_repo_name or re.search(r'[^a-zA-Z0-9_-]', new_repo_name) or len(new_repo_name) > 100:
+             logging.warning(f"Invalid repo name format received: {new_repo_name}")
+             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 IDLE and yield message
+             return package_workflow_outputs(
+                history=history, repo_id=repo_id, state=STATE_IDLE,
+                updated_preview=preview_html, updated_run=container_logs, updated_build=build_logs,
+                attempts=debug_attempts, app_desc=app_description, repo_name=repo_name, generated_code=generated_code,
+                use_grounding=use_grounding, current_gemini_key=current_gemini_key, current_gemini_model=current_gemini_model
+             )
+
+        history = add_bot_message(history, f"Acknowledged: '{message}'. Starting workflow to create Space `{hf_profile.username}/{new_repo_name}` for a '{new_app_desc}' app.")
+        logging.info(f"Transitioning to STATE_CREATING_SPACE for repo '{new_repo_name}' and description '{new_app_desc}'")
+        # Update state variables for the next step
+        return package_workflow_outputs(
+            history=history, repo_id=repo_id, state=STATE_CREATING_SPACE,
+            updated_preview=preview_html, updated_run=container_logs, updated_build=build_logs,
+            attempts=0, app_desc=new_app_desc, repo_name=new_repo_name, generated_code=None, # Reset attempts and generated_code
+            use_grounding=use_grounding, current_gemini_key=current_gemini_key, current_gemini_model=current_gemini_model
+        )
+
+    elif create_match:
+        logging.info("Simple create command received.")
+        new_repo_name = create_match.group(1).strip()
+        # Perform basic validation on repo name format
+        if not new_repo_name or re.search(r'[^a-zA-Z0-9_-]', new_repo_name) or len(new_repo_name) > 100:
+             logging.warning(f"Invalid repo name format received: {new_repo_name}")
+             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 IDLE and yield message
+             return package_workflow_outputs(
+                history=history, repo_id=repo_id, state=STATE_IDLE,
+                updated_preview=preview_html, updated_run=container_logs, updated_build=build_logs,
+                attempts=debug_attempts, app_desc=app_description, repo_name=repo_name, generated_code=generated_code,
+                use_grounding=use_grounding, current_gemini_key=current_gemini_key, current_gemini_model=current_gemini_model
+             )
+
+        history = add_bot_message(history, f"Acknowledged: '{message}'. Starting workflow to create Space `{hf_profile.username}/{new_repo_name}`.")
+        logging.info(f"Transitioning to STATE_CREATING_SPACE for repo '{new_repo_name}'")
+        return package_workflow_outputs(
+            history=history, repo_id=repo_id, state=STATE_CREATING_SPACE,
+            updated_preview=preview_html, updated_run=container_logs, updated_build=build_logs,
+            attempts=0, app_desc=app_description, repo_name=new_repo_name, generated_code=None, # Reset attempts and generated_code
+            use_grounding=use_grounding, current_gemini_key=current_gemini_key, current_gemini_model=current_gemini_model
+        )
+
+    elif "create" in message.lower() and not repo_id:
+        logging.info("Create command without name received.")
+        history = add_bot_message(history, "Okay, what should the Space be called? (e.g., `my-awesome-app`)")
+        logging.info("Transitioning to STATE_AWAITING_REPO_NAME")
+        return package_workflow_outputs(
+            history=history, repo_id=repo_id, state=STATE_AWAITING_REPO_NAME,
+            updated_preview=preview_html, updated_run=container_logs, updated_build=build_logs,
+            attempts=debug_attempts, app_desc=app_description, repo_name=repo_name, generated_code=generated_code,
+            use_grounding=use_grounding, current_gemini_key=current_gemini_key, current_gemini_model=current_gemini_model
+        )
+
+    else:
+        logging.info("Command not recognized in IDLE state.")
+        history = add_bot_message(history, "Command not recognized. Try 'generate me a gradio app called myapp', or 'reset'.")
+        # Stay in IDLE state
+        return package_workflow_outputs(
+            history=history, repo_id=repo_id, state=STATE_IDLE,
+            updated_preview=preview_html, updated_run=container_logs, updated_build=build_logs,
+            attempts=debug_attempts, app_desc=app_description, repo_name=repo_name, generated_code=generated_code,
+            use_grounding=use_grounding, current_gemini_key=current_gemini_key, current_gemini_model=current_gemini_model
+        )
+
+
+def handle_awaiting_repo_name(
+    message: str,
+    history: List[Dict[str, str]],
+    hf_profile: gr.OAuthProfile | None,
+    hf_token:   gr.OAuthToken   | None,
+    gemini_api_key: str | None,
+    gemini_model:   str | None,
+    repo_id:  str | None,
+    state: WorkflowState, # Should be STATE_AWAITING_REPO_NAME
+    space_sdk:      str,
+    preview_html:   str,
+    container_logs: str,
+    build_logs:     str,
+    debug_attempts: int,
+    app_description: str | None,
+    repo_name: str | None,
+    generated_code: str | None,
+    use_grounding: bool,
+    *args,
+    **kwargs
+) -> WorkflowOutputs:
+    """Handles logic when in the AWAITING_REPO_NAME state."""
+    logging.info(f"Handling STATE_AWAITING_REPO_NAME with message: {message[:50]}...")
+    current_gemini_key = gemini_api_key # Use the input vars directly
+    current_gemini_model = gemini_model
+
+    new_repo_name = message.strip()
+    # Basic validation for Hugging Face repo name format
+    # Allow letters, numbers, hyphens, underscores, max 100 chars (HF limit check)
+    if not new_repo_name or re.search(r'[^a-zA-Z0-9_-]', new_repo_name) or len(new_repo_name) > 100:
+        logging.warning(f"Invalid repo name format received while awaiting name: {new_repo_name}")
+        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
+        return package_workflow_outputs(
+            history=history, repo_id=repo_id, state=STATE_AWAITING_REPO_NAME,
+            updated_preview=preview_html, updated_run=container_logs, updated_build=build_logs,
+            attempts=debug_attempts, app_desc=app_description, repo_name=repo_name, generated_code=generated_code,
+            use_grounding=use_grounding, current_gemini_key=current_gemini_key, current_gemini_model=current_gemini_model
+        )
+    else:
+        history = add_bot_message(history, f"Using Space name `{new_repo_name}`. Creating Space `{hf_profile.username}/{new_repo_name}`...")
+        logging.info(f"Validated repo name '{new_repo_name}'. Transitioning to STATE_CREATING_SPACE.")
+        # Transition state to creation
+        return package_workflow_outputs(
+            history=history, repo_id=repo_id, state=STATE_CREATING_SPACE,
+            updated_preview=preview_html, updated_run=container_logs, updated_build=build_logs,
+            attempts=0, app_desc=app_description, repo_name=new_repo_name, generated_code=None, # Reset attempts and generated_code
+            use_grounding=use_grounding, current_gemini_key=current_gemini_key, current_gemini_model=current_gemini_model
+        )
+
+
+def handle_creating_space(
+    message: str,
+    history: List[Dict[str, str]],
+    hf_profile: gr.OAuthProfile | None,
+    hf_token:   gr.OAuthToken   | None,
+    gemini_api_key: str | None,
+    gemini_model:   str | None,
+    repo_id:  str | None,
+    state: WorkflowState, # Should be STATE_CREATING_SPACE
+    space_sdk:      str,
+    preview_html:   str,
+    container_logs: str,
+    build_logs:     str,
+    debug_attempts: int,
+    app_description: str | None,
+    repo_name: str | None,
+    generated_code: str | None,
+    use_grounding: bool,
+    *args,
+    **kwargs
+) -> WorkflowOutputs:
+    """Handles logic when in the CREATING_SPACE state."""
+    logging.info(f"Handling STATE_CREATING_SPACE for repo '{repo_name}'")
+    current_gemini_key = gemini_api_key # Use the input vars directly
+    current_gemini_model = gemini_model
+
+    # Ensure repo_name is available (it should have been set in the previous step)
+    if not repo_name:
+        logging.error("Internal error: Repo name missing in STATE_CREATING_SPACE. Resetting.")
+        history = add_bot_message(history, "Internal error: Repo name missing for creation. Resetting.")
+        # Reset state on error
+        return package_workflow_outputs(
+            history=history, repo_id=None, state=STATE_IDLE,
+            updated_preview="<p>Error creating space.</p>", updated_run="", updated_build="",
+            attempts=0, app_desc=None, repo_name=None, generated_code=None,
+            use_grounding=use_grounding, current_gemini_key=current_gemini_key, current_gemini_model=current_gemini_model
+        )
+    else:
+        try:
+            new_repo_id, iframe_html = create_space_action(repo_name, space_sdk, hf_profile, hf_token)
+            history = add_bot_message(history, f"✅ Space `{new_repo_id}` created. Click 'Send' to generate and upload code.")
+            logging.info(f"Space '{new_repo_id}' created. Transitioning to STATE_GENERATING_CODE.")
+            # Update state variables for the next step (generation)
+            return package_workflow_outputs(
+                history=history, repo_id=new_repo_id, state=STATE_GENERATING_CODE,
+                updated_preview=iframe_html, updated_run=container_logs, updated_build=build_logs,
+                attempts=debug_attempts, app_desc=app_description, repo_name=repo_name, generated_code=generated_code,
+                use_grounding=use_grounding, current_gemini_key=current_gemini_key, current_gemini_model=current_gemini_model
+            )
+        except RuntimeError as e: # Catch specific RuntimeErrors raised by actions
+            logging.error(f"Caught RuntimeError creating space: {e}")
+            history = add_bot_message(history, f"❌ Error creating space: {e}. Click 'reset'.")
+            # Reset state on error
+            return package_workflow_outputs(
+                history=history, repo_id=None, state=STATE_IDLE,
+                updated_preview="<p>Error creating space.</p>", updated_run="", updated_build="",
+                attempts=0, app_desc=None, repo_name=None, generated_code=None,
+                use_grounding=use_grounding, current_gemini_key=current_gemini_key, current_gemini_model=current_gemini_model
+            )
+
+def handle_generating_code(
+    message: str,
+    history: List[Dict[str, str]],
+    hf_profile: gr.OAuthProfile | None,
+    hf_token:   gr.OAuthToken   | None,
+    gemini_api_key: str | None,
+    gemini_model:   str | None,
+    repo_id:  str | None,
+    state: WorkflowState, # Should be STATE_GENERATING_CODE
+    space_sdk:      str,
+    preview_html:   str,
+    container_logs: str,
+    build_logs:     str,
+    debug_attempts: int,
+    app_description: str | None,
+    repo_name: str | None,
+    generated_code: str | None,
+    use_grounding: bool,
+    *args,
+    **kwargs
+) -> Generator[WorkflowOutputs, None, WorkflowOutputs]: # Specify return type as Generator
+    """Handles logic when in the GENERATING_CODE state."""
+    logging.info("Handling STATE_GENERATING_CODE")
+    current_gemini_key = gemini_api_key # Use the input vars directly
+    current_gemini_model = gemini_model
+
+    # Define the prompt for Gemini based on the app description or a default
+    prompt_desc = app_description if app_description else f'a simple {space_sdk} app'
+    prompt = f"""
+You are an AI assistant specializing in Hugging Face Spaces using the {space_sdk} SDK.
+Generate a full, single-file Python app based on:
+'{prompt_desc}'
+Ensure the code is runnable as `app.py` in a Hugging Face Space using the `{space_sdk}` SDK. Include necessary imports and setup.
+Return **only** the python code block for `app.py`. Do not include any extra text, explanations, or markdown outside the code block.
+"""
+    try:
+        history = add_bot_message(history, f"🧠 Generating `{prompt_desc}` `{space_sdk}` app (`app.py`) code with Gemini...")
+        if use_grounding:
+            history = add_bot_message(history, "(Using Grounding with Google Search)")
+        # Yield message before API call to show immediate feedback
+        # Use package_workflow_outputs to construct the tuple
+        yield package_workflow_outputs(
+            history=history, repo_id=repo_id, state=state,
+            updated_preview=preview_html, updated_run=container_logs, updated_build=build_logs,
+            attempts=debug_attempts, app_desc=app_description, repo_name=repo_name, generated_code=generated_code,
+            use_grounding=use_grounding, current_gemini_key=current_gemini_key, current_gemini_model=current_gemini_model
+        )
+
+        code = call_gemini(prompt, current_gemini_key, current_gemini_model, use_grounding=use_grounding)
+        code = code.strip()
+        # Clean up markdown code blocks
+        code = re.sub(r'^```python\s*', '', code, flags=re.MULTILINE).strip()
+        code = re.sub(r'^```\s*', '', code, flags=re.MULTILINE).strip() # Catch generic code blocks too
+        code = re.sub(r'\s*```$', '', code, flags=re.MULTILINE).strip()
+
+
+        if not code:
+            logging.warning("Gemini returned empty code.")
+            raise ValueError("Gemini returned empty code.")
+
+        history = add_bot_message(history, "✅ `app.py` code generated. Click 'Send' to upload.")
+        logging.info("Code generated. Transitioning to STATE_UPLOADING_APP_PY.")
+        # Transition state and store generated code
+        return package_workflow_outputs(
+            history=history, repo_id=repo_id, state=STATE_UPLOADING_APP_PY,
+            updated_preview=preview_html, updated_run=container_logs, updated_build=build_logs,
+            attempts=debug_attempts, app_desc=app_description, repo_name=repo_name, generated_code=code,
+            use_grounding=use_grounding, current_gemini_key=current_gemini_key, current_gemini_model=current_gemini_model
+        )
+
+    except RuntimeError as e: # Catch specific RuntimeErrors from call_gemini
+        logging.error(f"Caught RuntimeError generating code: {e}")
+        history = add_bot_message(history, f"❌ Error generating code: {e}. Click 'reset'.")
+        # Reset state on error
+        return package_workflow_outputs(
+            history=history, repo_id=None, state=STATE_IDLE,
+            updated_preview=preview_html, updated_run=container_logs, updated_build=build_logs,
+            attempts=0, app_desc=None, repo_name=None, generated_code=None,
+            use_grounding=use_grounding, current_gemini_key=current_gemini_key, current_gemini_model=current_gemini_model
+        )
+
+
+def handle_uploading_app_py(
+    message: str,
+    history: List[Dict[str, str]],
+    hf_profile: gr.OAuthProfile | None,
+    hf_token:   gr.OAuthToken   | None,
+    gemini_api_key: str | None,
+    gemini_model:   str | None,
+    repo_id:  str | None,
+    state: WorkflowState, # Should be STATE_UPLOADING_APP_PY
+    space_sdk:      str,
+    preview_html:   str,
+    container_logs: str,
+    build_logs:     str,
+    debug_attempts: int,
+    app_description: str | None,
+    repo_name: str | None,
+    generated_code: str | None, # This should hold the code to upload
+    use_grounding: bool,
+    *args,
+    **kwargs
+) -> Generator[WorkflowOutputs, None, WorkflowOutputs]: # Specify return type as Generator
+    """Handles logic when in the UPLOADING_APP_PY state."""
+    logging.info("Handling STATE_UPLOADING_APP_PY")
+    current_gemini_key = gemini_api_key # Use the input vars directly
+    current_gemini_model = gemini_model
+
+    # Retrieve the generated code from the state variable
+    code_to_upload = generated_code
+    if not code_to_upload:
+         logging.error("Internal error: No code to upload in STATE_UPLOADING_APP_PY. Resetting.")
+         history = add_bot_message(history, "Internal error: No code to upload. Resetting.")
+         return package_workflow_outputs(
+            history=history, repo_id=None, state=STATE_IDLE,
+            updated_preview=preview_html, updated_run=container_logs, updated_build=build_logs,
+            attempts=0, app_desc=None, repo_name=None, generated_code=None,
+            use_grounding=use_grounding, current_gemini_key=current_gemini_key, current_gemini_model=current_gemini_model
+         )
+    else:
+        history = add_bot_message(history, "☁️ Uploading `app.py`...")
+        # Yield message before upload
+        yield package_workflow_outputs(
+            history=history, repo_id=repo_id, state=state,
+            updated_preview=preview_html, updated_run=container_logs, updated_build=build_logs,
+            attempts=debug_attempts, app_desc=app_description, repo_name=repo_name, generated_code=generated_code,
+            use_grounding=use_grounding, current_gemini_key=current_gemini_key, current_gemini_model=current_gemini_model
+        )
+        try:
+            upload_file_to_space_action(io.StringIO(code_to_upload), "app.py", repo_id, hf_profile, hf_token)
+            history = add_bot_message(history, "✅ Uploaded `app.py`. Click 'Send' to generate requirements.")
+            logging.info("app.py uploaded. Transitioning to STATE_GENERATING_REQUIREMENTS.")
+            # Transition state, clear generated code after use
+            return package_workflow_outputs(
+                history=history, repo_id=repo_id, state=STATE_GENERATING_REQUIREMENTS,
+                updated_preview=preview_html, updated_run=container_logs, updated_build=build_logs,
+                attempts=debug_attempts, app_desc=app_description, repo_name=repo_name, generated_code=None,
+                use_grounding=use_grounding, current_gemini_key=current_gemini_key, current_gemini_model=current_gemini_model
+            )
+        except RuntimeError as e: # Catch specific RuntimeErrors from upload_file_to_space_action
+            logging.error(f"Caught RuntimeError uploading app.py: {e}")
+            history = add_bot_message(history, f"❌ Error uploading `app.py`: {e}. Click 'reset'.")
+            # Reset state on error
+            return package_workflow_outputs(
+                history=history, repo_id=None, state=STATE_IDLE,
+                updated_preview=preview_html, updated_run=container_logs, updated_build=build_logs,
+                attempts=0, app_desc=None, repo_name=None, generated_code=None,
+                use_grounding=use_grounding, current_gemini_key=current_gemini_key, current_gemini_model=current_gemini_model
+            )
+
+
+def handle_generating_requirements(
+    message: str,
+    history: List[Dict[str, str]],
+    hf_profile: gr.OAuthProfile | None,
+    hf_token:   gr.OAuthToken   | None,
+    gemini_api_key: str | None,
+    gemini_model:   str | None,
+    repo_id:  str | None,
+    state: WorkflowState, # Should be STATE_GENERATING_REQUIREMENTS
+    space_sdk:      str,
+    preview_html:   str,
+    container_logs: str,
+    build_logs:     str,
+    debug_attempts: int,
+    app_description: str | None,
+    repo_name: str | None,
+    generated_code: str | None,
+    use_grounding: bool,
+    *args,
+    **kwargs
+) -> Generator[WorkflowOutputs, None, WorkflowOutputs]: # Specify return type as Generator
+    """Handles logic when in the GENERATING_REQUIREMENTS state."""
+    logging.info("Handling STATE_GENERATING_REQUIREMENTS")
+    current_gemini_key = gemini_api_key # Use the input vars directly
+    current_gemini_model = gemini_model
+
+    history = add_bot_message(history, "📄 Generating `requirements.txt`...")
+    # Yield message before generating requirements
+    yield package_workflow_outputs(
+        history=history, repo_id=repo_id, state=state,
+        updated_preview=preview_html, updated_run=container_logs, updated_build=build_logs,
+        attempts=debug_attempts, app_desc=app_description, repo_name=repo_name, generated_code=generated_code,
+        use_grounding=use_grounding, current_gemini_key=current_gemini_key, current_gemini_model=current_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"]
+    # Add essential libraries regardless of description keywords or grounding
+    essential_libs = ["google-generativeai", "huggingface_hub"]
+    reqs_list.extend(essential_libs)
+
+    # Add common libraries if description suggests they might be needed
+    if app_description:
+        app_desc_lower = app_description.lower()
+        if "requests" in app_desc_lower or "api" in app_desc_lower:
+             reqs_list.append("requests")
+        if "image" in app_desc_lower or "upload" in app_desc_lower or "blur" in app_desc_lower or "vision" in app_desc_lower or "photo" in app_desc_lower:
+            reqs_list.append("Pillow")
+        if "numpy" in app_desc_lower: reqs_list.append("numpy")
+        if "pandas" in app_desc_lower or "dataframe" in app_desc_lower: reqs_list.append("pandas")
+        if any(lib in app_desc_lower for lib in ["scikit-image", "skimage", "cv2", "opencv"]):
+             reqs_list.extend(["scikit-image", "opencv-python"])
+        if any(lib in app_desc_lower for lib in ["transformer", "llama", "mistral", "bert", "gpt2"]):
+            reqs_list.append("transformers")
+        if any(lib in app_desc_lower for lib in ["torch", "pytorch", "tensorflow", "keras"]):
+            reqs_list.extend(["torch", "tensorflow"]) # Consider adding specific hardware versions if needed
+
+    # Use dict.fromkeys to get unique items while preserving insertion order (Python 3.7+)
+    reqs_list = list(dict.fromkeys(reqs_list))
+    # Sort alphabetically for cleaner requirements.txt
+    reqs_list.sort()
+
+    reqs_content = "\n".join(reqs_list) + "\n"
+
+    history = add_bot_message(history, "✅ `requirements.txt` generated. Click 'Send' to upload.")
+    logging.info("requirements.txt generated. Transitioning to STATE_UPLOADING_REQUIREMENTS.")
+    # Transition state and store requirements content
+    return package_workflow_outputs(
+        history=history, repo_id=repo_id, state=STATE_UPLOADING_REQUIREMENTS,
+        updated_preview=preview_html, updated_run=container_logs, updated_build=build_logs,
+        attempts=debug_attempts, app_desc=app_description, repo_name=repo_name, generated_code=reqs_content,
+        use_grounding=use_grounding, current_gemini_key=current_gemini_key, current_gemini_model=current_gemini_model
+    )
+
+
+def handle_uploading_requirements(
+    message: str,
+    history: List[Dict[str, str]],
+    hf_profile: gr.OAuthProfile | None,
+    hf_token:   gr.OAuthToken   | None,
+    gemini_api_key: str | None,
+    gemini_model:   str | None,
+    repo_id:  str | None,
+    state: WorkflowState, # Should be STATE_UPLOADING_REQUIREMENTS
+    space_sdk:      str,
+    preview_html:   str,
+    container_logs: str,
+    build_logs:     str,
+    debug_attempts: int,
+    app_description: str | None,
+    repo_name: str | None,
+    generated_code: str | None, # This should hold the requirements content
+    use_grounding: bool,
+    *args,
+    **kwargs
+) -> Generator[WorkflowOutputs, None, WorkflowOutputs]: # Specify return type as Generator
+    """Handles logic when in the UPLOADING_REQUIREMENTS state."""
+    logging.info("Handling STATE_UPLOADING_REQUIREMENTS")
+    current_gemini_key = gemini_api_key # Use the input vars directly
+    current_gemini_model = gemini_model
+
+    # Retrieve requirements content from state variable
+    reqs_content_to_upload = generated_code
+    if not reqs_content_to_upload:
+         logging.error("Internal error: No requirements content to upload in STATE_UPLOADING_REQUIREMENTS. Resetting.")
+         history = add_bot_message(history, "Internal error: No requirements content to upload. Resetting.")
+         return package_workflow_outputs(
+            history=history, repo_id=None, state=STATE_IDLE,
+            updated_preview=preview_html, updated_run=container_logs, updated_build=build_logs,
+            attempts=0, app_desc=None, repo_name=None, generated_code=None,
+            use_grounding=use_grounding, current_gemini_key=current_gemini_key, current_gemini_model=current_gemini_model
+         )
+    else:
+         history = add_bot_message(history, "☁️ Uploading `requirements.txt`...")
+         # Yield message before upload
+         yield package_workflow_outputs(
+            history=history, repo_id=repo_id, state=state,
+            updated_preview=preview_html, updated_run=container_logs, updated_build=build_logs,
+            attempts=debug_attempts, app_desc=app_description, repo_name=repo_name, generated_code=generated_code,
+            use_grounding=use_grounding, current_gemini_key=current_gemini_key, current_gemini_model=current_gemini_model
+         )
+         try:
+             # Perform requirements file upload
+             upload_file_to_space_action(io.StringIO(reqs_content_to_upload), "requirements.txt", repo_id, hf_profile, hf_token)
+             history = add_bot_message(history, "✅ Uploaded `requirements.txt`. Click 'Send' to generate README.")
+             logging.info("requirements.txt uploaded. Transitioning to STATE_GENERATING_README.")
+             # Transition state, clear generated code after use
+             return package_workflow_outputs(
+                history=history, repo_id=repo_id, state=STATE_GENERATING_README,
+                updated_preview=preview_html, updated_run=container_logs, updated_build=build_logs,
+                attempts=debug_attempts, app_desc=app_description, repo_name=repo_name, generated_code=None,
+                use_grounding=use_grounding, current_gemini_key=current_gemini_key, current_gemini_model=current_gemini_model
+             )
+         except RuntimeError as e: # Catch specific RuntimeErrors
+             logging.error(f"Caught RuntimeError uploading requirements.txt: {e}")
+             history = add_bot_message(history, f"❌ Error uploading `requirements.txt`: {e}. Click 'reset'.")
+             # Yield error message and reset state on failure
+             return package_workflow_outputs(
+                history=history, repo_id=None, state=STATE_IDLE,
+                updated_preview=preview_html, updated_run=container_logs, updated_build=build_logs,
+                attempts=0, app_desc=None, repo_name=None, generated_code=None,
+                use_grounding=use_grounding, current_gemini_key=current_gemini_key, current_gemini_model=current_gemini_model
+             )
+
+def handle_generating_readme(
+    message: str,
+    history: List[Dict[str, str]],
+    hf_profile: gr.OAuthProfile | None,
+    hf_token:   gr.OAuthToken   | None,
+    gemini_api_key: str | None,
+    gemini_model:   str | None,
+    repo_id:  str | None,
+    state: WorkflowState, # Should be STATE_GENERATING_README
+    space_sdk:      str,
+    preview_html:   str,
+    container_logs: str,
+    build_logs:     str,
+    debug_attempts: int,
+    app_description: str | None,
+    repo_name: str | None,
+    generated_code: str | None,
+    use_grounding: bool,
+    *args,
+    **kwargs
+) -> Generator[WorkflowOutputs, None, WorkflowOutputs]: # Specify return type as Generator
+    """Handles logic when in the GENERATING_README state."""
+    logging.info("Handling STATE_GENERATING_README")
+    current_gemini_key = gemini_api_key # Use the input vars directly
+    current_gemini_model = gemini_model
+
+    history = add_bot_message(history, "📝 Generating `README.md`...")
+    # Yield message before generating README
+    yield package_workflow_outputs(
+        history=history, repo_id=repo_id, state=state,
+        updated_preview=preview_html, updated_run=container_logs, updated_build=build_logs,
+        attempts=debug_attempts, app_desc=app_description, repo_name=repo_name, generated_code=generated_code,
+        use_grounding=use_grounding, current_gemini_key=current_gemini_key, current_gemini_model=current_gemini_model
+    )
+
+    # Generate simple README content with Space metadata header
+    readme_title = repo_name if repo_name else "My Awesome Space"
+    readme_description = app_description if app_description else f"This Hugging Face Space hosts an AI-generated {space_sdk} application."
+
+    readme_content = f"""---
+title: {readme_title}
+emoji: 🚀
+colorFrom: blue
+colorTo: yellow
+sdk: {space_sdk}
+app_file: app.py
+pinned: false
+---
+
+# {readme_title}
+
+{readme_description}
+
+This Space was automatically generated by an AI workflow using Google Gemini and Gradio.
+"""
+    history = add_bot_message(history, "✅ `README.md` generated. Click 'Send' to upload.")
+    logging.info("README.md generated. Transitioning to STATE_UPLOADING_README.")
+    # Transition state and store README content
+    return package_workflow_outputs(
+        history=history, repo_id=repo_id, state=STATE_UPLOADING_README,
+        updated_preview=preview_html, updated_run=container_logs, updated_build=build_logs,
+        attempts=debug_attempts, app_desc=app_description, repo_name=repo_name, generated_code=readme_content,
+        use_grounding=use_grounding, current_gemini_key=current_gemini_key, current_gemini_model=current_gemini_model
+    )
 
-    return gr.update(interactive=is_ready)
+
+def handle_uploading_readme(
+    message: str,
+    history: List[Dict[str, str]],
+    hf_profile: gr.OAuthProfile | None,
+    hf_token:   gr.OAuthToken   | None,
+    gemini_api_key: str | None,
+    gemini_model:   str | None,
+    repo_id:  str | None,
+    state: WorkflowState, # Should be STATE_UPLOADING_README
+    space_sdk:      str,
+    preview_html:   str,
+    container_logs: str,
+    build_logs:     str,
+    debug_attempts: int,
+    app_description: str | None,
+    repo_name: str | None,
+    generated_code: str | None, # This should hold the README content
+    use_grounding: bool,
+    *args,
+    **kwargs
+) -> Generator[WorkflowOutputs, None, WorkflowOutputs]: # Specify return type as Generator
+    """Handles logic when in the UPLOADING_README state."""
+    logging.info("Handling STATE_UPLOADING_README")
+    current_gemini_key = gemini_api_key # Use the input vars directly
+    current_gemini_model = gemini_model
+
+    # Retrieve README content from state variable
+    readme_content_to_upload = generated_code
+    if not readme_content_to_upload:
+         logging.error("Internal error: No README content to upload in STATE_UPLOADING_README. Resetting.")
+         history = add_bot_message(history, "Internal error: No README content to upload. Resetting.")
+         return package_workflow_outputs(
+            history=history, repo_id=None, state=STATE_IDLE,
+            updated_preview=preview_html, updated_run=container_logs, updated_build=build_logs,
+            attempts=0, app_desc=None, repo_name=None, generated_code=None,
+            use_grounding=use_grounding, current_gemini_key=current_gemini_key, current_gemini_model=current_gemini_model
+         )
+    else:
+         history = add_bot_message(history, "☁️ Uploading `README.md`...")
+         # Yield message before upload
+         yield package_workflow_outputs(
+            history=history, repo_id=repo_id, state=state,
+            updated_preview=preview_html, updated_run=container_logs, updated_build=build_logs,
+            attempts=debug_attempts, app_desc=app_description, repo_name=repo_name, generated_code=generated_code,
+            use_grounding=use_grounding, current_gemini_key=current_gemini_key, current_gemini_model=current_gemini_model
+         )
+         try:
+             # Perform README file upload
+             upload_file_to_space_action(io.StringIO(readme_content_to_upload), "README.md", repo_id, hf_profile, hf_token)
+             history = add_bot_message(history, "✅ Uploaded `README.md`. All files uploaded. Space is now building. Click 'Send' to check build logs.")
+             logging.info("README.md uploaded. Transitioning to STATE_CHECKING_LOGS_BUILD.")
+             # Transition to checking build logs, clear content after use
+             return package_workflow_outputs(
+                history=history, repo_id=repo_id, state=STATE_CHECKING_LOGS_BUILD,
+                updated_preview=preview_html, updated_run=container_logs, updated_build=build_logs,
+                attempts=debug_attempts, app_desc=app_description, repo_name=repo_name, generated_code=None,
+                use_grounding=use_grounding, current_gemini_key=current_gemini_key, current_gemini_model=current_gemini_model
+             )
+         except RuntimeError as e: # Catch specific RuntimeErrors
+             logging.error(f"Caught RuntimeError uploading README.md: {e}")
+             history = add_bot_message(history, f"❌ Error uploading `README.md`: {e}. Click 'reset'.")
+             # Yield error message and reset state on failure
+             return package_workflow_outputs(
+                history=history, repo_id=None, state=STATE_IDLE,
+                updated_preview=preview_html, updated_run=container_logs, updated_build=build_logs,
+                attempts=0, app_desc=None, repo_name=None, generated_code=None,
+                use_grounding=use_grounding, current_gemini_key=current_gemini_key, current_gemini_model=current_gemini_model
+             )
+
+def handle_checking_logs_build(
+    message: str,
+    history: List[Dict[str, str]],
+    hf_profile: gr.OAuthProfile | None,
+    hf_token:   gr.OAuthToken   | None,
+    gemini_api_key: str | None,
+    gemini_model:   str | None,
+    repo_id:  str | None,
+    state: WorkflowState, # Should be STATE_CHECKING_LOGS_BUILD
+    space_sdk:      str,
+    preview_html:   str,
+    container_logs: str, # Current UI value
+    build_logs:     str, # Current UI value
+    debug_attempts: int,
+    app_description: str | None,
+    repo_name: str | None,
+    generated_code: str | None,
+    use_grounding: bool,
+    *args,
+    **kwargs
+) -> Generator[WorkflowOutputs, None, WorkflowOutputs]: # Specify return type as Generator
+    """Handles logic when in the CHECKING_LOGS_BUILD state."""
+    logging.info(f"Handling STATE_CHECKING_LOGS_BUILD for repo '{repo_id}'")
+    current_gemini_key = gemini_api_key # Use the input vars directly
+    current_gemini_model = gemini_model
+
+    history = add_bot_message(history, "🔍 Fetching build logs...")
+    # Yield message before fetching logs (which includes a delay)
+    yield package_workflow_outputs(
+        history=history, repo_id=repo_id, state=state,
+        updated_preview=preview_html, updated_run=container_logs, updated_build=build_logs,
+        attempts=debug_attempts, app_desc=app_description, repo_name=repo_name, generated_code=generated_code,
+        use_grounding=use_grounding, current_gemini_key=current_gemini_key, current_gemini_model=current_gemini_model
+    )
+
+    # Fetch build logs from HF Space
+    build_logs_text = get_build_logs_action(repo_id, hf_profile, hf_token)
+    updated_build = build_logs_text # Update the logs display variable
+
+    # Simple check for common error indicators in logs (case-insensitive)
+    if "error" in updated_build.lower() or "exception" in updated_build.lower() or "build failed" in updated_build.lower():
+         logging.warning("Build logs indicate potential issues.")
+         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
+         logging.info("Build logs show issues. Transitioning to STATE_CHECKING_LOGS_RUN.")
+         # Yield updated state, logs, and variables
+         return package_workflow_outputs(
+            history=history, repo_id=repo_id, state=state,
+            updated_preview=preview_html, updated_run=container_logs, updated_build=updated_build, # Updated build logs
+            attempts=debug_attempts, app_desc=app_description, repo_name=repo_name, generated_code=generated_code,
+            use_grounding=use_grounding, current_gemini_key=current_gemini_key, current_gemini_model=current_gemini_model
+         )
+
+    else:
+         logging.info("Build logs appear clean.")
+         history = add_bot_message(history, "✅ Build logs fetched. Click 'Send' to check container logs.")
+         state = STATE_CHECKING_LOGS_RUN # Transition to next log check
+         logging.info("Transitioning to STATE_CHECKING_LOGS_RUN.")
+         # Yield updated state, logs, and variables
+         return package_workflow_outputs(
+            history=history, repo_id=repo_id, state=state,
+            updated_preview=preview_html, updated_run=container_logs, updated_build=updated_build, # Updated build logs
+            attempts=debug_attempts, app_desc=app_description, repo_name=repo_name, generated_code=generated_code,
+            use_grounding=use_grounding, current_gemini_key=current_gemini_key, current_gemini_model=current_gemini_model
+         )
+
+
+def handle_checking_logs_run(
+    message: str,
+    history: List[Dict[str, str]],
+    hf_profile: gr.OAuthProfile | None,
+    hf_token:   gr.OAuthToken   | None,
+    gemini_api_key: str | None,
+    gemini_model:   str | None,
+    repo_id:  str | None,
+    state: WorkflowState, # Should be STATE_CHECKING_LOGS_RUN
+    space_sdk:      str,
+    preview_html:   str,
+    container_logs: str, # Current UI value
+    build_logs:     str, # Current UI value
+    debug_attempts: int,
+    app_description: str | None,
+    repo_name: str | None,
+    generated_code: str | None,
+    use_grounding: bool,
+    *args,
+    **kwargs
+) -> Generator[WorkflowOutputs, None, WorkflowOutputs]: # Specify return type as Generator
+    """Handles logic when in the CHECKING_LOGS_RUN state."""
+    logging.info(f"Handling STATE_CHECKING_LOGS_RUN for repo '{repo_id}'")
+    current_gemini_key = gemini_api_key # Use the input vars directly
+    current_gemini_model = gemini_model
+
+    history = add_bot_message(history, "🔍 Fetching container logs...")
+    # Yield message before fetching logs (includes a delay)
+    yield package_workflow_outputs(
+        history=history, repo_id=repo_id, state=state,
+        updated_preview=preview_html, updated_run=container_logs, updated_build=build_logs,
+        attempts=debug_attempts, app_desc=app_description, repo_name=repo_name, generated_code=generated_code,
+        use_grounding=use_grounding, current_gemini_key=current_gemini_key, current_gemini_model=current_gemini_model
+    )
+
+    # Fetch container logs from HF Space
+    container_logs_text = get_container_logs_action(repo_id, hf_profile, hf_token)
+    updated_run = container_logs_text # Update the logs display variable
+
+    # Check for errors in run logs and if we have debug attempts left
+    if ("error" in updated_run.lower() or "exception" in updated_run.lower()) and debug_attempts < MAX_DEBUG_ATTEMPTS:
+         new_attempts = debug_attempts + 1 # Increment debug attempts counter
+         logging.warning(f"Errors detected in container logs. Attempting debug fix #{new_attempts}.")
+         history = add_bot_message(history, f"❌ Errors detected in container logs. Attempting debug fix #{new_attempts}/{MAX_DEBUG_ATTEMPTS}. Click 'Send' to proceed.")
+         state = STATE_DEBUGGING_CODE # Transition to the debugging state
+         logging.info("Transitioning to STATE_DEBUGGING_CODE.")
+         # Yield updated state, logs, attempts, and variables
+         return package_workflow_outputs(
+            history=history, repo_id=repo_id, state=state,
+            updated_preview=preview_html, updated_run=updated_run, updated_build=build_logs, # Updated run logs
+            attempts=new_attempts, app_desc=app_description, repo_name=repo_name, generated_code=generated_code,
+            use_grounding=use_grounding, current_gemini_key=current_gemini_key, current_gemini_model=current_gemini_model
+         )
+
+    elif ("error" in updated_run.lower() or "exception" in updated_run.lower()) and debug_attempts >= MAX_DEBUG_ATTEMPTS:
+         # Max debug attempts reached
+         logging.error(f"Errors detected in container logs. Max debug attempts ({MAX_DEBUG_ATTEMPTS}) reached.")
+         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
+         logging.info("Max debug attempts reached. Transitioning to STATE_COMPLETE.")
+         # Yield updated state, logs, attempts, and variables
+         return package_workflow_outputs(
+            history=history, repo_id=repo_id, state=state,
+            updated_preview=preview_html, updated_run=updated_run, updated_build=build_logs, # Updated run logs
+            attempts=debug_attempts, app_desc=app_description, repo_name=repo_name, generated_code=generated_code,
+            use_grounding=use_grounding, current_gemini_key=current_gemini_key, current_gemini_model=current_gemini_model
+         )
+
+    else:
+         # No significant errors found in logs, assume success
+         logging.info("No significant errors found in run logs.")
+         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
+         logging.info("Transitioning to STATE_COMPLETE.")
+         # Yield updated state, logs, attempts, and variables
+         return package_workflow_outputs(
+            history=history, repo_id=repo_id, state=state,
+            updated_preview=preview_html, updated_run=updated_run, updated_build=build_logs, # Updated run logs
+            attempts=debug_attempts, app_desc=app_description, repo_name=repo_name, generated_code=generated_code,
+            use_grounding=use_grounding, current_gemini_key=current_gemini_key, current_gemini_model=current_gemini_model
+         )
+
+
+def handle_debugging_code(
+    message: str,
+    history: List[Dict[str, str]],
+    hf_profile: gr.OAuthProfile | None,
+    hf_token:   gr.OAuthToken   | None,
+    gemini_api_key: str | None,
+    gemini_model:   str | None,
+    repo_id:  str | None,
+    state: WorkflowState, # Should be STATE_DEBUGGING_CODE
+    space_sdk:      str,
+    preview_html:   str,
+    container_logs: str, # Current UI value (contains logs to debug from)
+    build_logs:     str,
+    debug_attempts: int,
+    app_description: str | None,
+    repo_name: str | None,
+    generated_code: str | None,
+    use_grounding: bool,
+    *args,
+    **kwargs
+) -> Generator[WorkflowOutputs, None, WorkflowOutputs]: # Specify return type as Generator
+    """Handles logic when in the DEBUGGING_CODE state."""
+    logging.info(f"Handling STATE_DEBUGGING_CODE (attempt #{debug_attempts}) for repo '{repo_id}'")
+    current_gemini_key = gemini_api_key # Use the input vars directly
+    current_gemini_model = gemini_model
+
+    history = add_bot_message(history, f"🧠 Calling Gemini to generate fix based on logs...")
+    if use_grounding:
+         history = add_bot_message(history, "(Using Grounding with Google Search)")
+    # Yield message before Gemini API call
+    yield package_workflow_outputs(
+        history=history, repo_id=repo_id, state=state,
+        updated_preview=preview_html, updated_run=container_logs, updated_build=build_logs,
+        attempts=debug_attempts, app_desc=app_description, repo_name=repo_name, generated_code=generated_code,
+        use_grounding=use_grounding, current_gemini_key=current_gemini_key, current_gemini_model=current_gemini_model
+    )
+
+    # Construct prompt for Gemini including the container logs
+    debug_prompt = f"""
+You are debugging a {space_sdk} Space. The goal is to fix the code in `app.py` based on the container logs provided.
+
+Here are the container logs:
+Use code with caution.
+Python
+{container_logs}
+Generate the *complete, fixed* content for `app.py` based on these logs.
+Return **only** the python code block for app.py. Do not include any extra text, explanations, or markdown outside the code block.
+"""
+    try:
+         # Call Gemini to generate the corrected code, optionally using grounding
+         # Note: Grounding might be less effective for debugging based *only* on logs,
+         # but we include the option as requested.
+         # Use the current_gemini_key and current_gemini_model derived from state inputs
+         fix_code = call_gemini(debug_prompt, current_gemini_key, current_gemini_model, use_grounding=use_grounding)
+         fix_code = fix_code.strip()
+         # Clean up potential markdown formatting
+         fix_code = re.sub(r'^```python\s*', '', fix_code, flags=re.MULTILINE).strip()
+         fix_code = re.sub(r'^```\s*', '', fix_code, flags=re.MULTILINE).strip()
+         fix_code = re.sub(r'\s*```$', '', fix_code, flags=re.MULTILINE).strip()
+
+
+         if not fix_code:
+            logging.warning("Gemini returned empty fix code.")
+            raise ValueError("Gemini returned empty fix code.")
+
+         history = add_bot_message(history, "✅ Fix code generated. Click 'Send' to upload.")
+         logging.info("Fix code generated. Transitioning to STATE_UPLOADING_FIXED_APP_PY.")
+         # Transition to the upload state for the fix, store generated code
+         return package_workflow_outputs(
+            history=history, repo_id=repo_id, state=STATE_UPLOADING_FIXED_APP_PY,
+            updated_preview=preview_html, updated_run=container_logs, updated_build=build_logs,
+            attempts=debug_attempts, app_desc=app_description, repo_name=repo_name, generated_code=fix_code,
+            use_grounding=use_grounding, current_gemini_key=current_gemini_key, current_gemini_model=current_gemini_model
+         )
+
+    except RuntimeError as e: # Catch specific RuntimeErrors
+         logging.error(f"Caught RuntimeError generating debug code: {e}")
+         history = add_bot_message(history, f"❌ Error generating debug code: {e}. Click 'reset'.")
+         # Yield error message and reset state on failure
+         return package_workflow_outputs(
+            history=history, repo_id=None, state=STATE_IDLE,
+            updated_preview=preview_html, updated_run=container_logs, updated_build=build_logs,
+            attempts=0, app_desc=None, repo_name=None, generated_code=None,
+            use_grounding=use_grounding, current_gemini_key=current_gemini_key, current_gemini_model=current_gemini_model
+         )
+
+
+def handle_uploading_fixed_app_py(
+    message: str,
+    history: List[Dict[str, str]],
+    hf_profile: gr.OAuthProfile | None,
+    hf_token:   gr.OAuthToken   | None,
+    gemini_api_key: str | None,
+    gemini_model:   str | None,
+    repo_id:  str | None,
+    state: WorkflowState, # Should be STATE_UPLOADING_FIXED_APP_PY
+    space_sdk:      str,
+    preview_html:   str,
+    container_logs: str,
+    build_logs:     str,
+    debug_attempts: int,
+    app_description: str | None,
+    repo_name: str | None,
+    generated_code: str | None, # This should hold the fixed code
+    use_grounding: bool,
+    *args,
+    **kwargs
+) -> Generator[WorkflowOutputs, None, WorkflowOutputs]: # Specify return type as Generator
+    """Handles logic when in the UPLOADING_FIXED_APP_PY state."""
+    logging.info(f"Handling STATE_UPLOADING_FIXED_APP_PY for repo '{repo_id}'")
+    current_gemini_key = gemini_api_key # Use the input vars directly
+    current_gemini_model = gemini_model
+
+    # Retrieve the fixed code from the state variable
+    fixed_code_to_upload = generated_code
+    if not fixed_code_to_upload:
+         logging.error("Internal error: No fixed code available to upload in STATE_UPLOADING_FIXED_APP_PY. Resetting.")
+         history = add_bot_message(history, "Internal error: No fixed code available to upload. Resetting.")
+         return package_workflow_outputs(
+            history=history, repo_id=None, state=STATE_IDLE,
+            updated_preview=preview_html, updated_run=container_logs, updated_build=build_logs,
+            attempts=0, app_desc=None, repo_name=None, generated_code=None,
+            use_grounding=use_grounding, current_gemini_key=current_gemini_key, current_gemini_model=current_gemini_model
+         )
+    else:
+         history = add_bot_message(history, "☁️ Uploading fixed `app.py`...")
+         # Yield message before upload
+         yield package_workflow_outputs(
+            history=history, repo_id=repo_id, state=state,
+            updated_preview=preview_html, updated_run=container_logs, updated_build=build_logs,
+            attempts=debug_attempts, app_desc=app_description, repo_name=repo_name, generated_code=generated_code,
+            use_grounding=use_grounding, current_gemini_key=current_gemini_key, current_gemini_model=current_gemini_model
+         )
+
+         try:
+             # Perform the upload of the fixed app.py
+             upload_file_to_space_action(io.StringIO(fixed_code_to_upload), "app.py", repo_id, hf_profile, hf_token)
+             history = add_bot_message(history, "✅ Fixed `app.py` uploaded. Space will rebuild. Click 'Send' to check logs again.")
+             state = STATE_CHECKING_LOGS_RUN # Go back to checking run logs after uploading the fix
+             logging.info("Fixed app.py uploaded. Transitioning to STATE_CHECKING_LOGS_RUN.")
+             # Transition state, clear code after use
+             return package_workflow_outputs(
+                history=history, repo_id=repo_id, state=state,
+                updated_preview=preview_html, updated_run=container_logs, updated_build=build_logs,
+                attempts=debug_attempts, app_desc=app_description, repo_name=repo_name, generated_code=None,
+                use_grounding=use_grounding, current_gemini_key=current_gemini_key, current_gemini_model=current_gemini_model
+             )
+         except RuntimeError as e: # Catch specific RuntimeErrors
+             logging.error(f"Caught RuntimeError uploading fixed app.py: {e}")
+             history = add_bot_message(history, f"❌ Error uploading fixed `app.py`: {e}. Click 'reset'.")
+             # Yield error message and reset state on failure
+             return package_workflow_outputs(
+                history=history, repo_id=None, state=STATE_IDLE,
+                updated_preview=preview_html, updated_run=container_logs, updated_build=build_logs,
+                attempts=0, app_desc=None, repo_name=None, generated_code=None,
+                use_grounding=use_grounding, current_gemini_key=current_gemini_key, current_gemini_model=current_gemini_model
+             )
+
+def handle_complete(
+    message: str, # User might type something in COMPLETE state
+    history: List[Dict[str, str]],
+    hf_profile: gr.OAuthProfile | None,
+    hf_token:   gr.OAuthToken   | None,
+    gemini_api_key: str | None,
+    gemini_model:   str | None,
+    repo_id:  str | None,
+    state: WorkflowState, # Should be STATE_COMPLETE
+    space_sdk:      str,
+    preview_html:   str,
+    container_logs: str,
+    build_logs:     str,
+    debug_attempts: int,
+    app_description: str | None,
+    repo_name: str | None,
+    generated_code: str | None,
+    use_grounding: bool,
+    *args,
+    **kwargs
+) -> WorkflowOutputs:
+    """Handles logic when in the COMPLETE state."""
+    logging.info("Handling STATE_COMPLETE")
+    current_gemini_key = gemini_api_key # Use the input vars directly
+    current_gemini_model = gemini_model
+
+    # If the user types something in the complete state, maybe interpret it?
+    # For now, we'll just stay in COMPLETE unless they type 'reset'.
+    if "reset" in message.lower():
+        logging.info("Reset command received in COMPLETE state.")
+        history = add_bot_message(history, "Workflow reset.")
+        # Reset relevant states and UI outputs
+        return package_workflow_outputs(
+            history=history, repo_id=None, state=STATE_IDLE,
+            updated_preview="<p>No Space created yet.</p>", updated_run="", updated_build="",
+            attempts=0, app_desc=None, repo_name=None, generated_code=None,
+            use_grounding=use_grounding, current_gemini_key=current_gemini_key, current_gemini_model=current_gemini_model
+        )
+    else:
+         # Stay in COMPLETE state
+         history = add_bot_message(history, "Workflow is complete. Type 'reset' to start a new project.")
+         return package_workflow_outputs(
+            history=history, repo_id=repo_id, state=STATE_COMPLETE,
+            updated_preview=preview_html, updated_run=container_logs, updated_build=build_logs,
+            attempts=debug_attempts, app_desc=app_description, repo_name=repo_name, generated_code=generated_code,
+            use_grounding=use_grounding, current_gemini_key=current_gemini_key, current_gemini_model=current_gemini_model
+         )
+
+
+# --- Dispatch Table ---
+STATE_HANDLERS: Dict[WorkflowState, Any] = { # Use Any for type hint simplicity here
+    STATE_IDLE: handle_idle,
+    STATE_AWAITING_REPO_NAME: handle_awaiting_repo_name,
+    STATE_CREATING_SPACE: handle_creating_space,
+    STATE_GENERATING_CODE: handle_generating_code,
+    STATE_UPLOADING_APP_PY: handle_uploading_app_py,
+    STATE_GENERATING_REQUIREMENTS: handle_generating_requirements,
+    STATE_UPLOADING_REQUIREMENTS: handle_uploading_requirements,
+    STATE_GENERATING_README: handle_generating_readme,
+    STATE_UPLOADING_README: handle_uploading_readme,
+    STATE_CHECKING_LOGS_BUILD: handle_checking_logs_build,
+    STATE_CHECKING_LOGS_RUN: handle_checking_logs_run,
+    STATE_DEBUGGING_CODE: handle_debugging_code,
+    STATE_UPLOADING_FIXED_APP_PY: handle_uploading_fixed_app_py,
+    STATE_COMPLETE: handle_complete,
+}
 
 
 # This is the main generator function for the workflow, triggered by the 'Send' button
 # Inputs and Outputs list must match exactly. The generator receives values from the inputs list.
 def ai_workflow_chat(
     message: str,
-    history: list[dict],
+    history: List[Dict[str, str]],
     hf_profile: gr.OAuthProfile | None,
     hf_token:   gr.OAuthToken   | None,
     # Pass gemini_api_key and gemini_model as inputs - these come from the State variables
     gemini_api_key_state: str | None,
     gemini_model_state:   str | None,
     repo_id_state:  str | None,
-    workflow_state: str,
+    workflow_state: WorkflowState, # Use the Literal type hint
     space_sdk:      str,
     # NOTE: UI component values are passed *by value* to the generator
     preview_html:   str, # Value from iframe HTML
@@ -305,814 +1431,267 @@ def ai_workflow_chat(
     # Accept any extra args/kwargs passed by Gradio, common for generators
     *args,
     **kwargs
-) -> tuple[
-    list[dict],       # 0: Updated chat history (for chatbot)
-    str | None,       # 1: Updated repo_id (for repo_id state)
-    str,              # 2: Updated workflow state (for workflow state)
-    str,              # 3: Updated iframe HTML (for iframe UI component)
-    str,              # 4: Updated container logs (for run_txt UI component)
-    str,              # 5: Updated build logs (for build_txt UI component)
-    int,              # 6: Updated debug attempts count (for debug_attempts state)
-    str | None,       # 7: Updated app description (for app_description state)
-    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)
-    str | None,       # 11: Explicitly yield gemini_api_key_state
-    str | None,       # 12: Explicitly yield gemini_model_state
-]:
+) -> Any: # Use Any because it yields multiple times before returning the final value (None in this case)
     """
     Generator function to handle the AI workflow state machine.
     Each 'yield' pauses execution and sends values to update Gradio outputs/state.
     """
-    # Unpack state variables from Gradio State components passed as inputs
+    # Unpack state variables and UI values from Gradio inputs
     repo_id = repo_id_state
     state = workflow_state
     attempts = debug_attempts_state
     app_desc = app_description_state
     repo_name = repo_name_state
     generated_code = generated_code_state
-    use_grounding = use_grounding_state # Unpack grounding state
-    # Use the input parameters for Gemini key/model directly in the generator
+    use_grounding = use_grounding_state
     current_gemini_key = gemini_api_key_state
     current_gemini_model = gemini_model_state
 
-    # --- START DEBUGGING ai_workflow_chat inputs ---
-    print("\n--- ai_workflow_chat START (Inputs received) ---")
-    print(f"  message: {message}")
-    print(f"  history len: {len(history)}")
-    print(f"  hf_profile: {hf_profile is not None}")
-    print(f"  hf_token: {hf_token is not None}")
-    api_key_display = current_gemini_key[:5] if isinstance(current_gemini_key, str) and current_gemini_key else ('Empty String' if isinstance(current_gemini_key, str) and current_gemini_key == "" else 'None')
-    print(f"  gemini_api_key_state: Value starts with '{api_key_display}'")
-    print(f"  gemini_model_state: {current_gemini_model}")
-    print(f"  repo_id_state: {repo_id_state}") # Check value here
-    print(f"  workflow_state: {workflow_state}")
-    print(f"  space_sdk: {space_sdk}")
-    print(f"  use_grounding_state: {use_grounding_state}")
-    print(f"  debug_attempts_state: {debug_attempts_state}")
-    print(f"  app_description_state: {app_description_state}")
-    print(f"  repo_name_state: {repo_name_state}")
-    print(f"  generated_code_state: {'Present' if generated_code_state is not None else 'None'}")
-    print(f"  *args (unexpected by generator): {args}") # Added debug for unexpected args
-    print(f"  **kwargs (unexpected by generator): {kwargs}") # Added debug for unexpected kwargs
-    print("--- END DEBUGGING ai_workflow_chat inputs ---\n")
-
-
-    # Keep copies of potentially updated UI elements passed as inputs to update them later
-    # These are the *current values* of the UI components as of the button click
-    updated_preview = preview_html
-    updated_build = build_logs
-    updated_run = container_logs
+    logging.info(f"ai_workflow_chat generator started. State: {state}, Message: {message[:50]}...")
+    # Log all inputs for debugging if needed
+    # logging.debug(f"ai_workflow_chat inputs: {locals()}")
 
     # Add the user's message to the chat history immediately
     user_message_entry = {"role": "user", "content": message}
-    # Add username if logged in (optional, but nice)
     if hf_profile and hf_profile.username:
          user_message_entry["name"] = hf_profile.username
     history.append(user_message_entry)
+    logging.debug("User message added to history.")
+
 
     # 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
-    # Include gemini_api_key_state and gemini_model_state in the yield tuple
-    # The yielded tuple must match the send_btn.click outputs list exactly.
-    yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-           attempts, app_desc, repo_name, generated_code, use_grounding,
-           current_gemini_key, current_gemini_model) # Explicitly pass back current state values
+    yield package_workflow_outputs(
+        history=history, repo_id=repo_id, state=state,
+        updated_preview=preview_html, updated_run=container_logs, updated_build=build_logs,
+        attempts=attempts, app_desc=app_desc, repo_name=repo_name, generated_code=generated_code,
+        use_grounding=use_grounding, current_gemini_key=current_gemini_key, current_gemini_model=current_gemini_model
+    )
 
     try:
-        # --- State Machine Logic based on the current 'state' variable ---
-
-        # Although button interactivity prevents reaching here without key/model,
-        # the checks remain as a safeguard for the workflow logic itself.
-        # Use the local variables derived from state inputs (current_gemini_key, current_gemini_model)
-        if not (hf_profile and hf_token):
-            history = add_bot_message(history, "Workflow paused: Please log in to Hugging Face first.")
-            # Re-yield state to update chat and keep current state values
-            yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                   attempts, app_desc, repo_name, generated_code, use_grounding,
-                   current_gemini_key, current_gemini_model)
-            return # Stop workflow execution for this click
-
-        if not (isinstance(current_gemini_key, str) and current_gemini_key != "" and current_gemini_model):
-             history = add_bot_message(history, "Workflow cannot start: Please ensure your Gemini API key is entered and a model is selected.")
-             # Re-yield state to update chat and keep current state values
-             yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                    attempts, app_desc, repo_name, generated_code, use_grounding,
-                    current_gemini_key, current_gemini_model) # Correctly yielding back the received (likely empty/None) Gemini states
-             return # Stop workflow execution for this click
-
-
-        if state == STATE_IDLE:
-            # Look for specific commands in the user's message
-            reset_match = "reset" in message.lower()
-            # Capture app description AND repo name using regex
-            generate_match = re.search(r'generate (?:me )?(?:a|an) (.+) app called (\w+)', message, re.I)
-            # Capture repo name for a simple 'create space' command
-            create_match = re.search(r'create (?:a|an)? space called (\w+)', message, re.I)
-
-            if reset_match:
-                history = add_bot_message(history, "Workflow reset.")
-                # Reset relevant states and UI outputs, passing through current Gemini state
-                yield (history, None, STATE_IDLE, "<p>No Space created yet.</p>", "", "", 0,
-                       None, None, None, False, current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
-
-            elif generate_match:
-                # User requested generation with description and name
-                new_app_desc = generate_match.group(1).strip() # Capture description part
-                new_repo_name = generate_match.group(2).strip() # Capture name part
-                history = add_bot_message(history, f"Acknowledged: '{message}'. Starting workflow to create Space `{hf_profile.username}/{new_repo_name}` for a '{new_app_desc}' app.")
-                # Update state variables for the next step (creation)
-                state = STATE_CREATING_SPACE
-                repo_name = new_repo_name
-                app_desc = new_app_desc
-                # Yield updated state variables, passing others through
-                yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                       attempts, app_desc, repo_name, generated_code, use_grounding,
-                       current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
-
-            elif create_match:
-                 # User requested simple space creation with a name
-                 new_repo_name = create_match.group(1).strip()
-                 history = add_bot_message(history, f"Acknowledged: '{message}'. Starting workflow to create Space `{hf_profile.username}/{new_repo_name}`.")
-                 state = STATE_CREATING_SPACE # Transition state to creation
-                 repo_name = new_repo_name # Store the validated repo name
-                 # Yield updated state variables, passing others through
-                 yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                        attempts, app_desc, repo_name, generated_code, use_grounding,
-                        current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
-
-            elif "create" in message.lower() and not repo_id:
-                # User wants to create but didn't specify a name yet
-                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 state, passing others through
-                yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                       attempts, app_desc, repo_name, generated_code, use_grounding,
-                       current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
-
-            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 current state, passing others through
-                yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                       attempts, app_desc, repo_name, generated_code, use_grounding,
-                       current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
-
-
-        elif state == STATE_AWAITING_REPO_NAME:
-             # User's message is expected to be the repo name
-             new_repo_name = message.strip()
-             # Basic validation for Hugging Face repo name format
-             # Allow letters, numbers, hyphens, underscores, max 100 chars (HF limit check)
-             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)
-                 yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                        attempts, app_desc, repo_name, generated_code, use_grounding,
-                        current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
-
-             else:
-                 history = add_bot_message(history, f"Using Space name `{new_repo_name}`. Creating Space `{hf_profile.username}/{new_repo_name}`...")
-                 state = STATE_CREATING_SPACE # Transition state to creation
-                 repo_name = new_repo_name # Store the validated repo name
-                 # Yield updated state variables, passing others through
-                 yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                        attempts, app_desc, repo_name, generated_code, use_grounding,
-                        current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
-
-
-        elif state == STATE_CREATING_SPACE:
-             # Ensure repo_name is available (it should have been set in the previous step)
-             if not repo_name:
-                 history = add_bot_message(history, "Internal error: Repo name missing for creation. Resetting.")
-                 # Reset state on error, passing through current Gemini state
-                 yield (history, None, STATE_IDLE, "<p>Error creating space.</p>", "", "", 0,
-                       None, None, None, use_grounding, current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
-             else:
-                 try:
-                     new_repo_id, iframe_html = create_space_action(repo_name, space_sdk, hf_profile, hf_token)
-                     updated_preview = iframe_html
-                     repo_id = new_repo_id # Update repo_id state variable
-                     history = add_bot_message(history, f"✅ Space `{repo_id}` created. Click 'Send' to generate and upload code.")
-                     state = STATE_GENERATING_CODE
-                     # Yield updated state variables, passing others through
-                     yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                            attempts, app_desc, repo_name, generated_code, use_grounding,
-                            current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
-                 except Exception as e:
-                     history = add_bot_message(history, f"❌ Error creating space: {e}. Click 'reset'.")
-                     # Reset state on error, passing through current Gemini state
-                     yield (history, None, STATE_IDLE, "<p>Error creating space.</p>", "", "", 0,
-                           None, None, None, use_grounding, current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
-
-
-        elif state == STATE_GENERATING_CODE:
-             # Define the prompt for Gemini based on the app description or a default
-             prompt_desc = app_desc if app_desc else f'a simple {space_sdk} app'
-             prompt = f"""
-You are an AI assistant specializing in Hugging Face Spaces using the {space_sdk} SDK.
-Generate a full, single-file Python app based on:
-'{prompt_desc}'
-Ensure the code is runnable as `app.py` in a Hugging Face Space using the `{space_sdk}` SDK. Include necessary imports and setup.
-Return **only** the python code block for `app.py`. Do not include any extra text, explanations, or markdown outside the code block.
-"""
-             try:
-                 history = add_bot_message(history, f"🧠 Generating `{prompt_desc}` `{space_sdk}` app (`app.py`) code with Gemini...")
-                 if use_grounding:
-                     history = add_bot_message(history, "(Using Grounding with Google Search)")
-                 # Yield message before API call
-                 yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                        attempts, app_desc, repo_name, generated_code, use_grounding,
-                        current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
-
-                 code = call_gemini(prompt, current_gemini_key, current_gemini_model, use_grounding=use_grounding)
-                 code = code.strip()
-                 # Clean up markdown
-                 if code.startswith("```python"): code = code[len("```python"):].strip()
-                 if code.startswith("```"): code = code[len("```"):].strip()
-                 if code.endswith("```"): code = code[:-len("```")].strip()
-
-                 if not code: raise ValueError("Gemini returned empty code.")
-
-                 history = add_bot_message(history, "✅ `app.py` code generated. Click 'Send' to upload.")
-                 state = STATE_UPLOADING_APP_PY
-                 generated_code = code
-                 # Yield updated state variables
-                 yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                        attempts, app_desc, repo_name, generated_code, use_grounding,
-                        current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
-
-             except Exception as e:
-                 history = add_bot_message(history, f"❌ Error generating code: {e}. Click 'reset'.")
-                 # Reset state on error, passing through current Gemini state
-                 yield (history, None, STATE_IDLE, updated_preview, updated_run, updated_build, 0,
-                       None, None, None, use_grounding, current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
-
-
-        elif state == STATE_UPLOADING_APP_PY:
-             # Retrieve the generated code from the state variable
-             code_to_upload = generated_code
-             if not code_to_upload:
-                  history = add_bot_message(history, "Internal error: No code to upload. Resetting.")
-                  yield (history, None, STATE_IDLE, updated_preview, updated_run, updated_build, 0,
-                        None, None, None, use_grounding, current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
-             else:
-                 history = add_bot_message(history, "☁️ Uploading `app.py`...")
-                 # Yield message before upload
-                 yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                        attempts, app_desc, repo_name, generated_code, use_grounding,
-                        current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
-                 try:
-                     upload_file_to_space_action(io.StringIO(code_to_upload), "app.py", repo_id, hf_profile, hf_token)
-                     history = add_bot_message(history, "✅ Uploaded `app.py`. Click 'Send' to generate requirements.")
-                     state = STATE_GENERATING_REQUIREMENTS
-                     generated_code = None
-                     # Yield updated state variables
-                     yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                            attempts, app_desc, repo_name, generated_code, use_grounding,
-                            current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
-                 except Exception as e:
-                     history = add_bot_message(history, f"❌ Error uploading `app.py`: {e}. Click 'reset'.")
-                     # Reset state on error, passing through current Gemini state
-                     yield (history, None, STATE_IDLE, updated_preview, updated_run, updated_build, 0,
-                           None, None, None, use_grounding, current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
-
-
-        elif state == STATE_GENERATING_REQUIREMENTS:
-             history = add_bot_message(history, "📄 Generating `requirements.txt`...")
-             # Yield message before generating requirements
-             yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                    attempts, app_desc, repo_name, generated_code, use_grounding,
-                    current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
-
-             # Logic to determine required packages based on SDK and keywords in the app description
-             reqs_list = ["gradio"] if space_sdk == "gradio" else ["streamlit"]
-             # Add essential libraries regardless of description keywords or grounding
-             essential_libs = ["google-generativeai", "huggingface_hub"]
-             # Only add if Gemini is actually needed for the app (determined by description or if key is present)
-             # If we are here, key and model are available based on STATE_IDLE checks
-             reqs_list.extend(essential_libs)
-
-             # Add common libraries if description suggests they might be needed
-             if app_desc:
-                 app_desc_lower = app_desc.lower()
-                 if "requests" in app_desc_lower or "api" in app_desc_lower:
-                      reqs_list.append("requests")
-                 # Image processing libraries
-                 if "image" in app_desc_lower or "upload" in app_desc_lower or "blur" in app_desc_lower or "vision" in app_desc_lower or "photo" in app_desc_lower:
-                     reqs_list.append("Pillow")
-                 if "numpy" in app_desc_lower: reqs_list.append("numpy")
-                 if "pandas" in app_desc_lower or "dataframe" in app_desc_lower: reqs_list.append("pandas")
-                 # Add scikit-image and opencv if image processing is heavily implied
-                 if any(lib in app_desc_lower for lib in ["scikit-image", "skimage", "cv2", "opencv"]):
-                      reqs_list.extend(["scikit-image", "opencv-python"]) # Note: opencv-python for pip
-                 # Add transformers if large models are implied
-                 if any(lib in app_desc_lower for lib in ["transformer", "llama", "mistral", "bert", "gpt2"]):
-                     reqs_list.append("transformers")
-                 # Add torch or tensorflow if deep learning frameworks are implied
-                 if any(lib in app_desc_lower for lib in ["torch", "pytorch", "tensorflow", "keras"]):
-                     reqs_list.extend(["torch", "tensorflow"]) # Users might need specific versions, but this is a start
-
-             # Use dict.fromkeys to get unique items while preserving insertion order (Python 3.7+)
-             reqs_list = list(dict.fromkeys(reqs_list))
-             # Sort alphabetically for cleaner requirements.txt
-             reqs_list.sort()
-
-             reqs_content = "\n".join(reqs_list) + "\n"
-
-             history = add_bot_message(history, "✅ `requirements.txt` generated. Click 'Send' to upload.")
-             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)
-             yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                    attempts, app_desc, repo_name, generated_code, use_grounding,
-                    current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
-
-
-        elif state == STATE_UPLOADING_REQUIREMENTS:
-            # Retrieve requirements content from state variable
-            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.")
-                 yield (history, None, STATE_IDLE, updated_preview, updated_run, updated_build, 0,
-                        None, None, None, use_grounding, current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
-            else:
-                 history = add_bot_message(history, "☁️ Uploading `requirements.txt`...")
-                 # Yield message before upload (pass UI outputs and states through)
-                 yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                        attempts, app_desc, repo_name, generated_code, use_grounding,
-                        current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
-                 try:
-                     # Perform requirements file upload
-                     upload_file_to_space_action(io.StringIO(reqs_content_to_upload), "requirements.txt", repo_id, hf_profile, hf_token)
-                     history = add_bot_message(history, "✅ Uploaded `requirements.txt`. Click 'Send' to generate README.")
-                     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)
-                     yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                            attempts, app_desc, repo_name, generated_code, use_grounding,
-                            current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
-                 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
-                     yield (history, None, STATE_IDLE, updated_preview, updated_run, updated_build, 0,
-                           None, None, None, use_grounding, current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
-
-        elif state == STATE_GENERATING_README:
-            history = add_bot_message(history, "📝 Generating `README.md`...")
-            # Yield message before generating README (pass UI outputs and states through)
-            yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                   attempts, app_desc, repo_name, generated_code, use_grounding,
-                   current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
-
-            # Generate simple README content with Space metadata header
-            readme_title = repo_name if repo_name else "My Awesome Space"
-            readme_description = app_desc if app_desc else f"This Hugging Face Space hosts an AI-generated {space_sdk} application."
-
-            readme_content = f"""---
-title: {readme_title}
-emoji: 🚀
-colorFrom: blue
-colorTo: yellow
-sdk: {space_sdk}
-app_file: app.py
-pinned: false
----
-
-# {readme_title}
-
-{readme_description}
-
-This Space was automatically generated by an AI workflow using Google Gemini and Gradio.
-""" # Added Space metadata header and slightly improved content
-
-            history = add_bot_message(history, "✅ `README.md` generated. Click 'Send' to upload.")
-            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)
-            yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                   attempts, app_desc, repo_name, generated_code, use_grounding,
-                   current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
-
-
-        elif state == STATE_UPLOADING_README:
-            # Retrieve README content from state variable
-            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.")
-                 yield (history, None, STATE_IDLE, updated_preview, updated_run, updated_build, 0,
-                        None, None, None, use_grounding, current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
+        # --- State Machine Logic using Dispatch Table ---
+        handler = STATE_HANDLERS.get(state)
+
+        if handler:
+            logging.debug(f"Invoking handler for state: {state}")
+            # Call the state handler function, passing all necessary data
+            # Need to pass *all* inputs to the handler function
+            handler_output = handler(
+                message=message, history=history,
+                hf_profile=hf_profile, hf_token=hf_token,
+                gemini_api_key=current_gemini_key, gemini_model=current_gemini_model, # Pass current values
+                repo_id=repo_id, state=state, space_sdk=space_sdk,
+                preview_html=preview_html, container_logs=container_logs, build_logs=build_logs, # Pass current UI values
+                debug_attempts=attempts, app_description=app_desc, repo_name=repo_name, generated_code=generated_code, # Pass current state values
+                use_grounding=use_grounding
+            )
+            # The handler might yield intermediate updates (e.g., "Generating...")
+            if isinstance(handler_output, Generator):
+                 # If the handler is also a generator, yield from it
+                 logging.debug("Handler is a generator, yielding from it.")
+                 yield from handler_output
             else:
-                 history = add_bot_message(history, "☁️ Uploading `README.md`...")
-                 # Yield message before upload (pass UI outputs and states through)
-                 yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                        attempts, app_desc, repo_name, generated_code, use_grounding,
-                        current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
-                 try:
-                     # Perform README file upload
-                     upload_file_to_space_action(io.StringIO(readme_content_to_upload), "README.md", repo_id, hf_profile, hf_token)
-                     history = add_bot_message(history, "✅ Uploaded `README.md`. All files uploaded. Space is now building. Click 'Send' to check build logs.")
-                     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)
-                     yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                            attempts, app_desc, repo_name, generated_code, use_grounding,
-                            current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
-                 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
-                     yield (history, None, STATE_IDLE, updated_preview, updated_run, updated_build, 0,
-                           None, None, None, use_grounding, current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
-
-        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)
-             yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                    attempts, app_desc, repo_name, generated_code, use_grounding,
-                    current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
-
-             # Fetch build logs from HF Space
-             build_logs_text = get_build_logs_action(repo_id, hf_profile, hf_token)
-             updated_build = build_logs_text # Update the logs display variable
-
-             # Simple check for common error indicators in logs (case-insensitive)
-             if "error" in updated_build.lower() or "exception" in updated_build.lower() or "build failed" in updated_build.lower():
-                  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
-                  yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                         attempts, app_desc, repo_name, generated_code, use_grounding,
-                         current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
-
-             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
-                  yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                         attempts, app_desc, repo_name, generated_code, use_grounding,
-                         current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
-
-
-        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)
-             yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                    attempts, app_desc, repo_name, generated_code, use_grounding,
-                    current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
-
-             # Fetch container logs from HF Space
-             container_logs_text = get_container_logs_action(repo_id, hf_profile, hf_token)
-             updated_run = container_logs_text # Update the logs display variable
-
-             # Check for errors in run logs and if we have debug attempts left
-             if ("error" in updated_run.lower() or "exception" in updated_run.lower()) and attempts < MAX_DEBUG_ATTEMPTS:
-                  attempts += 1 # Increment debug attempts counter
-                  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
-                  yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                         attempts, app_desc, repo_name, generated_code, use_grounding,
-                         current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
-
-             elif ("error" in updated_run.lower() or "exception" in updated_run.lower()) and attempts >= MAX_DEBUG_ATTEMPTS:
-                  # Max debug attempts reached
-                  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
-                  yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                         attempts, app_desc, repo_name, generated_code, use_grounding,
-                         current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
-
-             else:
-                  # No significant errors found in logs, assume success
-                  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
-                  yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                         attempts, app_desc, repo_name, generated_code, use_grounding,
-                         current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
-
-
-        elif state == STATE_DEBUGGING_CODE:
-             history = add_bot_message(history, f"🧠 Calling Gemini to generate fix based on logs...")
-             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)
-             yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                    attempts, app_desc, repo_name, generated_code, use_grounding,
-                    current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
-
-             # Construct prompt for Gemini including the container logs
-             debug_prompt = f"""
-You are debugging a {space_sdk} Space. The goal is to fix the code in `app.py` based on the container logs provided.
-
-Here are the container logs:
-Use code with caution.
-Python
-{updated_run}
-Generate the *complete, fixed* content for `app.py` based on these logs.
-Return **only** the python code block for app.py. Do not include any extra text, explanations, or markdown outside the code block.
-"""
-             try:
-                  # Call Gemini to generate the corrected code, optionally using grounding
-                  # Note: Grounding might be less effective for debugging based *only* on logs,
-                  # but we include the option as requested.
-                  # Use the current_gemini_key and current_gemini_model derived from state inputs
-                  fix_code = call_gemini(debug_prompt, current_gemini_key, current_gemini_model, use_grounding=use_grounding)
-                  fix_code = fix_code.strip()
-                  # Clean up potential markdown formatting
-                  if fix_code.startswith("```python"):
-                      fix_code = fix_code[len("```python"):].strip()
-                  if fix_code.startswith("```"):
-                       fix_code = fix_code[len("```"):].strip()
-                  if fix_code.endswith("```"):
-                      fix_code = fix_code[:-len("```")].strip()
-
-                  if not fix_code:
-                     raise ValueError("Gemini returned empty fix code.")
-
-                  history = add_bot_message(history, "✅ Fix code generated. Click 'Send' to upload.")
-                  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)
-                  yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                         attempts, app_desc, repo_name, generated_code, use_grounding,
-                         current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
-
-                  # After uploading a fix, we usually go back to checking logs
-                  state = STATE_CHECKING_LOGS_RUN # Correct transition after upload fix
-
-             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
-                  yield (history, None, STATE_IDLE, updated_preview, updated_run, updated_build, 0,
-                         None, None, None, use_grounding,
-                         current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
-
-        elif state == STATE_UPLOADING_FIXED_APP_PY:
-             # Retrieve the fixed code from the state variable
-             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.")
-                  yield (history, None, STATE_IDLE, updated_preview, updated_run, updated_build, 0,
-                        None, None, None, use_grounding, current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
-             else:
-                  history = add_bot_message(history, "☁️ Uploading fixed `app.py`...")
-                  # Yield message before upload (pass UI outputs and states through)
-                  yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                         attempts, app_desc, repo_name, generated_code, use_grounding,
-                         current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
-
-                  try:
-                      # Perform the upload of the fixed app.py
-                      upload_file_to_space_action(io.StringIO(fixed_code_to_upload), "app.py", repo_id, hf_profile, hf_token)
-                      history = add_bot_message(history, "✅ Fixed `app.py` uploaded. Space will rebuild. Click 'Send' to check logs again.")
-                      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)
-                      yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                             attempts, app_desc, repo_name, generated_code, use_grounding,
-                             current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
-
-                  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
-                      yield (history, None, STATE_IDLE, updated_preview, updated_run, updated_build, 0,
-                            None, None, None, use_grounding, current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
-
-        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.
-             yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                    attempts, app_desc, repo_name, generated_code, use_grounding,
-                    current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
+                 # If the handler returned the final tuple for this step, yield it
+                 logging.debug("Handler returned final output tuple, yielding it.")
+                 yield handler_output
+
+        else:
+            logging.error(f"No handler found for state: {state}. Resetting.")
+            # Fallback for unknown state
+            history = add_bot_message(history, f"Internal error: Unknown state `{state}`. Resetting.")
+            yield package_workflow_outputs(
+                history=history, repo_id=None, state=STATE_IDLE,
+                updated_preview="<p>Error: Unknown state.</p>", updated_run="", updated_build="",
+                attempts=0, app_desc=None, repo_name=None, generated_code=None,
+                use_grounding=use_grounding, current_gemini_key=current_gemini_key, current_gemini_model=current_gemini_model
+            )
 
 
     except Exception as e:
         # This catches any unexpected errors that occur within any state's logic
+        # Specific errors from helper functions (like RuntimeError) should ideally be caught in handlers,
+        # but this is a safety net.
         error_message = f"Workflow step failed unexpectedly ({state}): {e}. Click 'Send' to re-attempt this step or 'reset'."
         history = add_bot_message(history, error_message)
-        print(f"Critical Error in state {state}: {e}") # Log the error for debugging purposes
+        logging.exception(f"Critical Error caught in ai_workflow_chat generator for state {state}") # Log with traceback
+
         # On unexpected error, reset to IDLE, but pass through the current Gemini state
-        yield (history, None, STATE_IDLE, updated_preview, updated_run, updated_build, 0,
-               None, None, None, use_grounding,
-               current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
+        yield package_workflow_outputs(
+            history=history, repo_id=None, state=STATE_IDLE,
+            updated_preview=preview_html, updated_run=container_logs, updated_build=build_logs, # Keep existing UI logs
+            attempts=0, app_desc=None, repo_name=None, generated_code=None, # Reset project-specific states
+            use_grounding=use_grounding, current_gemini_key=current_gemini_key, current_gemini_model=current_gemini_model # Pass through Gemini states
+        )
 
 
 # --- Build the Gradio UI ---
 
 with gr.Blocks(title="AI-Powered HF Space App Builder") as ai_builder_tab:
     # Gradio State variables - these persist their values across user interactions (clicks)
-    # Define these first as they might be used in default values for components
     hf_profile   = gr.State(None)
     hf_token     = gr.State(None)
-    # Initialize gemini_api_key_state to empty string
     gemini_api_key_state   = gr.State("")  # start with no key
-    # Initialize gemini_model_state with the default model key
     gemini_model_state     = gr.State(DEFAULT_GEMINI_MODEL) # Default selected model
 
     repo_id      = gr.State(None) # Stores the ID of the created Space
-    workflow     = gr.State(STATE_IDLE) # Stores the current state of the AI workflow
+    workflow     = gr.State(STATE_IDLE, live=True) # Stores the current state, live update for status_text
     sdk_state    = gr.State("gradio") # Stores the selected Space SDK (Gradio or Streamlit)
     debug_attempts = gr.State(0) # Counter for how many debugging attempts have been made
     app_description = gr.State(None) # Stores the user's initial description of the desired app
     repo_name_state = gr.State(None) # Stores the chosen repository name for the Space
     generated_code_state = gr.State(None) # Temporary storage for generated file content (app.py, reqs, README)
-    # New State variable for grounding checkbox
     use_grounding_state = gr.State(False)
 
 
     with gr.Row():
-        # Sidebar column for inputs and status displays
         with gr.Column(scale=1, min_width=300):
             gr.Markdown("## Hugging Face Login")
-            # Define login_status before it's used in login_btn.click outputs
             login_status = gr.Markdown("*Not logged in.*")
-            # Hugging Face Login Button
             login_btn    = gr.LoginButton(variant="huggingface")
 
             gr.Markdown("## Google AI Studio / Gemini")
-            # Define gemini_input and gemini_status before they are used in change handlers
-            # Blank out textbox on load and update info text
             gemini_input  = gr.Textbox(
-                label="Your Google AI Studio API Key", # Changed label
-                type="password", # Hides input for security
+                label="Your Google AI Studio API Key",
+                type="password",
                 interactive=True,
-                value="", # Don't pre-fill from env var
-                info="Enter your own key here" # Updated info text
+                value="",
+                info="Enter your own key here"
             )
-            gemini_status = gr.Markdown("") # Display Gemini configuration status
+            gemini_status = gr.Markdown("")
 
-            # Define model_selector before it's used in its change handler
             model_selector = gr.Radio(
-                # Use the list of choices generated from the GEMINI_MODELS dictionary
                 choices=GEMINI_MODEL_CHOICES,
-                value=DEFAULT_GEMINI_MODEL, # Default selection using the key
+                value=DEFAULT_GEMINI_MODEL,
                 label="Select model",
                 interactive=True
             )
 
-            # Add a markdown field to display the model description
-            # Initialize with the description of the default model
             model_description_text = gr.Markdown(get_model_description(DEFAULT_GEMINI_MODEL))
 
-
-            # Define grounding checkbox before its change handler
             use_grounding_checkbox = gr.Checkbox(
                 label="Enable Grounding with Google Search",
-                value=False, # Default to off
+                value=False,
                 interactive=True,
                 info="Use Google Search results to inform Gemini's response (may improve factuality)."
             )
 
             gr.Markdown("## Space SDK")
-            # Define sdk_selector before its change handler
             sdk_selector = gr.Radio(choices=["gradio","streamlit"], value="gradio", label="Template SDK", interactive=True)
 
             gr.Markdown("## Workflow Status")
-            # Define status_text and repo_id_text before they are used in change handlers
             status_text = gr.Textbox(label="Current State", value=STATE_IDLE, interactive=False)
             repo_id_text = gr.Textbox(label="Current Space ID", value="None", interactive=False)
 
 
-        # Main content area column
         with gr.Column(scale=3):
-            # Define chatbot, user_input, send_btn before send_btn.click
             chatbot    = gr.Chatbot(type='messages', label="AI Workflow Chat")
             user_input = gr.Textbox(placeholder="Type your message…", interactive=True)
-            # Define send_btn here, BEFORE it's used in send_button_update_output
-            send_btn   = gr.Button("Send", interactive=False) # Start disabled by default
+            send_btn   = gr.Button("Send", interactive=False)
 
 
-            # Define iframe, build_txt, run_txt after send_btn
-            # These are UI components, NOT State variables
-            iframe    = gr.HTML("<p>No Space created yet.</p>") # HTML element for the Space iframe
-            build_txt = gr.Textbox(label="Build Logs", lines=10, interactive=False, value="", max_lines=20) # Set max_lines for scrollability
-            run_txt   = gr.Textbox(label="Container Logs", lines=10, interactive=False, value="", max_lines=20) # Set max_lines for scrollability
+            iframe    = gr.HTML("<p>No Space created yet.</p>")
+            build_txt = gr.Textbox(label="Build Logs", lines=10, interactive=False, value="", max_lines=20)
+            run_txt   = gr.Textbox(label="Container Logs", lines=10, interactive=False, value="", max_lines=20)
 
 
-    # --- Define Event Handlers and Chains AFTER all components and required lists are defined ---
+    # --- Define Event Handlers and Chains ---
 
-    # Define the inputs used for checking prerequisites (These are State components)
-    send_button_interactive_binding_inputs = [
-        hf_profile,
-        hf_token,
-        gemini_api_key_state,
-        gemini_model_state
+    # List of prerequisite State components for the send button logic
+    prerequisite_states_for_button = [
+        hf_profile, hf_token, gemini_api_key_state, gemini_model_state, workflow # Add workflow state
     ]
-    # Define the output for updating the send button interactivity
-    send_button_update_output = [send_btn]
-
-    # Lambda to pass only the required 4 state inputs to check_send_button_ready,
-    # ignoring the first argument (the changed value) passed by state.change
-    lambda_check_button_from_state_change = lambda changed_value, s1, s2, s3, s4: check_send_button_ready(s1, s2, s3, s4)
 
-    # Lambda to pass only the required 4 state inputs to check_send_button_ready,
-    # ignoring the previous step's output passed by .then
-    lambda_check_button_from_then = lambda prev_output, s1, s2, s3, s4: check_send_button_ready(s1, s2, s3, s4)
-
-
-    # Trigger check_send_button_ready whenever any prerequisite state changes
-    # Use the specific lambda which expects (changed_value, *input_values)
-    hf_profile.change(
-        lambda_check_button_from_state_change,
-        inputs=send_button_interactive_binding_inputs, # Pass all 4 prerequisite states
-        outputs=send_button_update_output, # Update only the send button
-    )
-    hf_token.change(
-        lambda_check_button_from_state_change,
-        inputs=send_button_interactive_binding_inputs,
-        outputs=send_button_update_output,
-    )
-    # When gemini_api_key_state changes (updated by gemini_input.change), check button readiness
-    # The .then() chain attached to gemini_input.change updates the state *first*, then triggers this change handler.
-    # This change handler receives (new_state_value, *inputs_values).
-    gemini_api_key_state.change(
-        lambda_check_button_from_state_change,
-        inputs=send_button_interactive_binding_inputs,
-        outputs=send_button_update_output,
-    )
-    # When gemini_model_state changes (updated by model_selector.change), check button readiness
-    # This change handler receives (new_state_value, *inputs_values).
-    gemini_model_state.change(
-        lambda_check_button_from_state_change,
-        inputs=send_button_interactive_binding_inputs,
-        outputs=send_button_update_output,
-    )
+    # Use the pattern suggested in the feedback: wire each dependency change to the same handler
+    for state_comp in prerequisite_states_for_button:
+        state_comp.change(
+            # Lambda function receives the new value of the changed component first,
+            # followed by the values of the components in 'inputs'.
+            # We pass ALL prerequisite states (including the one that changed) to the lambda's inputs.
+            # The lambda then passes the *explicitly listed* input values to the target function, check_send_button_ready.
+            # This avoids relying on the order of implicit args.
+            lambda *args, states=prerequisite_states_for_button: check_send_button_ready(
+                states[0], states[1], states[2], states[3], states[4] # Pass values from the 'states' list closure
+            ),
+            inputs=prerequisite_states_for_button, # Pass all required states
+            outputs=[send_btn], # Update only the send button
+        )
+        # Add a debug log to confirm wiring (optional debug)
+        # logging.debug(f"Wired {state_comp.label}.change to check_send_button_ready.")
 
 
     # Handle login button click: Update profile/token state -> Their .change handlers trigger check_send_button_ready
     login_btn.click(
          # Lambda takes the LoginButton output (profile, token tuple) which is 2 args: (profile, token)
          lambda profile, token: (profile, token),
-         inputs=[login_btn], # Pass the LoginButton itself to get its output (its output is a tuple of profile, token)
-         outputs=[hf_profile, hf_token] # Update state variables
-    )
+         inputs=[login_btn],
+         outputs=[hf_profile, hf_token]
+    ) # The .change handlers on hf_profile and hf_token will trigger check_send_button_ready
 
     # Handle Gemini Key Input change: Update key state -> Configure Gemini status
     gemini_input.change(
-        # Lambda receives the new value of gemini_input (1 arg) because inputs=None (implied)
+        # Lambda receives the new value of gemini_input (1 arg) because inputs=[gemini_input]
         lambda new_key_value: new_key_value,
-        inputs=None, # Only need the new value of the changed component
-        outputs=[gemini_api_key_state] # This output becomes the implicit first arg for the next .then() in this chain
+        inputs=[gemini_input], # Explicitly pass the changed component for clarity
+        outputs=[gemini_api_key_state] # This output updates the state
     ).then(
         # Configure Gemini using the updated state variables
-        # Lambda receives (prev_output, api_key_val, model_name_val) where prev_output is the new key value
-        # It needs to pass api_key_val and model_name_val to configure_gemini.
-        # Note: In this specific chain, prev_output IS the new key value, so we pass args[0] and args[2].
-        # This is fragile! A better approach is to rely *only* on explicit inputs in .then.
-        # Let's use the values from the *states* directly in the .then inputs.
+        # Lambda receives (prev_output, api_key_val_from_state, model_name_val_from_state)
+        # The prev_output is the new key value from the previous step's output (gemini_api_key_state)
+        # We use the explicit inputs instead of prev_output for robustness.
         lambda prev_output, api_key_val_from_state, model_name_val_from_state: configure_gemini(api_key_val_from_state, model_name_val_from_state),
         inputs=[gemini_api_key_state, gemini_model_state], # Explicitly pass the required states
         outputs=[gemini_status] # Update Gemini status display.
-    )
-    # The .then chain continues from the outputs of the configure_gemini call, but no further .then is attached here.
-    # The gemini_api_key_state.change handler (defined above) will pick up the state change.
-
+    ) # The gemini_api_key_state.change handler (wired in the loop above) handles button updates.
 
     # Handle Gemini Model Selector change: Update model state -> Update description -> Configure Gemini status
     model_selector.change(
-        # Lambda receives the new value of model_selector (1 arg)
+        # Lambda receives the new value of model_selector (1 arg) because inputs=[model_selector]
         lambda new_model_name: new_model_name,
-        inputs=None, # Only need the new value of the changed component
-        outputs=[gemini_model_state] # This output becomes the implicit first arg for the next .then()
+        inputs=[model_selector], # Explicitly pass the changed component for clarity
+        outputs=[gemini_model_state] # This output updates the state
     ).then(
         # Update the model description display
-        # Lambda receives (prev_output, model_name_val) where prev_output is the new model name
-        # It needs to pass model_name_val to get_model_description.
-        # Again, fragile. Rely on explicit inputs.
+        # Lambda receives (prev_output, model_name_val_from_state)
+        # The prev_output is the new model name from the previous step's output (gemini_model_state)
+        # We use the explicit inputs instead of prev_output for robustness.
         lambda prev_output, model_name_val_from_state: get_model_description(model_name_val_from_state),
         inputs=[gemini_model_state], # Get the new state value
         outputs=[model_description_text] # Update description UI.
     ).then(
         # Configure Gemini using the updated state variables
-        # Lambda receives (prev_output, api_key_val, model_name_val) where prev_output is the description text
-        # It needs to pass api_key_val and model_name_val to configure_gemini.
+        # Lambda receives (prev_output, api_key_val_from_state, model_name_val_from_state)
+        # The prev_output is the description text from the previous step.
+        # We use the explicit inputs instead of prev_output for robustness.
         lambda prev_output, api_key_val_from_state, model_name_val_from_state: configure_gemini(api_key_val_from_state, model_name_val_from_state),
         inputs=[gemini_api_key_state, gemini_model_state], # Explicitly pass the required states
         outputs=[gemini_status] # Update Gemini status display.
-    )
-    # The .then chain continues from the outputs of the configure_gemini call, handled by gemini_model_state.change handler above
+    ) # The gemini_model_state.change handler (wired in the loop above) handles button updates.
 
 
     # Handle Grounding checkbox change: update grounding state
     use_grounding_checkbox.change(
-        lambda v: v, inputs=use_grounding_checkbox, outputs=use_grounding_state
+        lambda v: v, inputs=[use_grounding_checkbox], outputs=[use_grounding_state] # Use lists for inputs/outputs
     )
 
     # Handle SDK selector change: update sdk state
     sdk_selector.change(
-        lambda s: s, inputs=sdk_selector, outputs=sdk_state
+        lambda s: s, inputs=[sdk_selector], outputs=[sdk_state] # Use lists for inputs/outputs
     )
 
     # Link Workflow State variable change to UI status display
-    # Lambda receives the new state value (1 arg) because inputs=None
-    workflow.change(lambda new_state_value: new_state_value, inputs=None, outputs=status_text)
+    workflow.change(
+        lambda new_state_value: new_state_value,
+        inputs=[workflow], # Use lists for inputs
+        outputs=[status_text] # Use lists for outputs
+    )
 
     # Link Repo ID State variable change to UI status display
-    # Lambda receives the new state value (1 arg) because inputs=None
-    repo_id.change(lambda new_repo_id_value: new_repo_id_value if new_repo_id_value else "None", inputs=None, outputs=repo_id_text)
+    repo_id.change(
+        lambda new_repo_id_value: new_repo_id_value if new_repo_id_value else "None",
+        inputs=[repo_id], # Use lists for inputs
+        outputs=[repo_id_text] # Use lists for outputs
+    )
 
 
     # The main event handler for the Send button (generator)
@@ -1122,69 +1701,74 @@ with gr.Blocks(title="AI-Powered HF Space App Builder") as ai_builder_tab:
     # Ensure inputs and outputs match the ai_workflow_chat signature and yield tuple EXACTLY.
     # This call is direct, not in a .then() chain, so it does NOT receive a prev_output arg.
     # It receives args only from the inputs list.
+    send_btn_inputs = [
+        user_input, chatbot, # UI component inputs (message, current chat history)
+        hf_profile, hf_token, # HF State variables
+        gemini_api_key_state, gemini_model_state, # Gemini State variables
+        repo_id, workflow, sdk_state, # Workflow State variables
+        iframe, run_txt, build_txt, # UI component inputs (current values)
+        debug_attempts, app_description, repo_name_state, generated_code_state, # Other State variables
+        use_grounding_state # Grounding state input
+    ]
+    send_btn_outputs = [
+        chatbot, # Updates Chatbot
+        repo_id, workflow, # Updates State variables (repo_id, workflow)
+        iframe, run_txt, build_txt, # Updates UI components (iframe, logs)
+        debug_attempts, app_description, repo_name_state, generated_code_state, # Updates other State variables
+        use_grounding_state, # Updates Grounding state
+        gemini_api_key_state, gemini_model_state # Updates Gemini State variables - these are passed through the generator
+    ]
+
     send_btn.click(
-        ai_workflow_chat, # The generator function to run
-        inputs=[
-            user_input, chatbot, # UI component inputs (message, current chat history)
-            hf_profile, hf_token, # HF State variables
-            gemini_api_key_state, gemini_model_state, # Gemini State variables
-            repo_id, workflow, sdk_state, # Workflow State variables
-            iframe, run_txt, build_txt, # UI component inputs (current values)
-            debug_attempts, app_description, repo_name_state, generated_code_state, # Other State variables
-            use_grounding_state # Grounding state input
-        ],
-        outputs=[
-            chatbot, # Updates Chatbot
-            repo_id, workflow, # Updates State variables (repo_id, workflow)
-            iframe, run_txt, build_txt, # Updates UI components (iframe, logs)
-            debug_attempts, app_description, repo_name_state, generated_code_state, # Updates other State variables
-            use_grounding_state, # Updates Grounding state
-            gemini_api_key_state, gemini_model_state # Updates Gemini State variables
-        ]
+        ai_workflow_chat,
+        inputs=send_btn_inputs,
+        outputs=send_btn_outputs
     ).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
-         lambda: gr.update(value=""),
+         lambda: gr.Textbox.update(value=""), # Use specific component update
          inputs=None,
-         outputs=user_input # Update the user input textbox
+         outputs=[user_input] # Use lists for outputs
     )
 
 
-    # --- Initial Load Event Chain (Defined INSIDE gr.Blocks, AFTER components and required lists are defined) ---
+    # --- Initial Load Event Chain ---
     # This chain runs once when the app loads
     ai_builder_tab.load(
         # Action 1: Show profile (loads cached login if available)
-        # show_profile receives no implicit args from load when inputs=None is used for load itself
-        # (If inputs was not None for load, it would receive those input values first).
-        # Let's use inputs=None for load for simplicity. show_profile expects 1 arg (profile).
-        # How does Gradio handle this for load? The load *function* takes no inputs, but the *event* passes the initial state values.
-        # Let's add a lambda wrapper to be safe.
-        lambda *args: show_profile(args[0] if args else None), # Expects 0 or 1 args, needs 1 for show_profile
+        # Lambda receives args corresponding to load's inputs. Load has no explicit inputs here.
+        # However, Gradio *does* pass the initial values of all components/states defined *before* the load event.
+        # The most robust way is to pass the specific state needed.
+        lambda initial_profile: show_profile(initial_profile),
         inputs=[hf_profile], # Pass the initial profile state value
-        outputs=login_status # Updates UI. This output becomes the implicit first arg for the next .then()
+        outputs=[login_status] # Updates UI. Use lists for outputs. This output becomes prev_output for the next .then()
     ).then(
         # Action 2: Configure Gemini using initial state
         # Lambda receives (prev_output, api_key_val, model_name_val)
+        # prev_output is the string from show_profile. Use explicit inputs.
         lambda prev_output, api_key_val_from_state, model_name_val_from_state: configure_gemini(api_key_val_from_state, model_name_val_from_state),
         inputs=[gemini_api_key_state, gemini_model_state], # Explicitly pass the required states
-        outputs=[gemini_status] # Update Gemini status display. This output becomes the implicit first arg for the next .then()
+        outputs=[gemini_status] # Update Gemini status display. Use lists for outputs.
     ).then(
         # Action 3: After initial load checks, update the button state based on initial states
         # Lambda receives (prev_output, *prereq_state_values)
-        lambda prev_output, s1, s2, s3, s4: check_send_button_ready(s1, s2, s3, s4),
-        inputs=send_button_interactive_binding_inputs, # Pass all 4 prerequisite states
-        outputs=send_button_update_output, # Update the send button. This output becomes implicit first arg for next .then()
+        # prev_output is the string from configure_gemini. Use explicit inputs.
+        lambda prev_output, p1, p2, p3, p4, p5: check_send_button_ready(p1, p2, p3, p4, p5), # Match check_send_button_ready signature
+        inputs=prerequisite_states_for_button, # Pass all 5 prerequisite states
+        outputs=[send_btn], # Update the send button. Use lists for outputs.
     ).then(
         # Action 4: Update the model description text based on the default selected model
         # Lambda receives (prev_output, model_name_val)
+        # prev_output is the gr.Button.update object. Use explicit input.
         lambda prev_output, model_name_val_from_state: get_model_description(model_name_val_from_state),
         inputs=[gemini_model_state], # Get the default model name from state
-        outputs=[model_description_text] # Update description UI. This output becomes implicit first arg for next .then()
+        outputs=[model_description_text] # Update description UI. Use lists for outputs.
     ).then(
         # Action 5: Add the initial welcome message to the chat history
         # Lambda receives (prev_output)
+        # prev_output is the description text.
         lambda prev_output: greet(),
         inputs=None, # Greet takes no explicit inputs
-        outputs=chatbot # Updates the chatbot display
+        outputs=[chatbot] # Updates the chatbot display. Use lists for outputs.
     )
 
 
@@ -1193,22 +1777,11 @@ with gr.Blocks(title="AI-Powered HF Space App Builder") as ai_builder_tab:
 
 
 if __name__ == "__main__":
-    # Optional: Configure retries for huggingface_hub requests to make them more robust
-    # from requests.adapters import HTTPAdapter
-    # from urllib3.util.retry import Retry
-    # retry_strategy = Retry(total=5, backoff_factor=1, status_forcelist=[429, 500, 502, 503, 504]) # Define retry strategy for specific HTTP codes
-    # adapter = HTTPAdapter(max_retries=retry_strategy)
-    # session = get_session() # Get the session object used internally by huggingface_hub
-    # session.mount("http://", adapter)
-    # session.mount("https://", adapter)
-
     # Optional: Configure Gradio settings using environment variables
-    # Set max upload size (e.g., 100MB) for files like app.py
     os.environ["GRADIO_MAX_FILE_SIZE"] = "100MB"
-    # Optional: Set a local temporary directory for Gradio uploads
     os.environ["GRADIO_TEMP_DIR"] = "./tmp"
-    os.makedirs(os.environ["GRADIO_TEMP_DIR"], exist_ok=True) # Ensure the directory exists
+    os.makedirs(os.environ["GRADIO_TEMP_DIR"], exist_ok=True)
 
+    logging.info("Starting Gradio app...")
     # Launch the Gradio UI
-    # The Gradio launch call blocks execution.
     ai_builder_tab.launch()
\ No newline at end of file