Update app.py

app.py

@@ -1,20 +1,133 @@
-# ... (existing imports and utility functions like show_profile, list_private_models,
-#      get_sdk_version, extract_code, classify_errors, _get_space_jwt,
-#      fetch_logs, check_iframe remain the same)
+import re
+import json
+import time
+import requests
+import importlib.metadata
+import gradio as gr
+import os # Needed for writing files
+from huggingface_hub import (
+    create_repo, upload_file, list_models, constants
+)
+from huggingface_hub.utils import build_hf_headers, get_session, hf_raise_for_status
+from google import genai
+from google.genai.types import Tool, GenerateContentConfig, GoogleSearch
+
+# --- USER INFO & MODEL LISTING ---
+
+def show_profile(profile: gr.OAuthProfile | None) -> str:
+    return f"✅ Logged in as **{profile.username}**" if profile else "*Not logged in.*"
+
+def list_private_models(
+    profile: gr.OAuthProfile | None,
+    oauth_token: gr.OAuthToken | None
+) -> str:
+    if not profile or not oauth_token:
+        return "Please log in to see your models."
+    try:
+        models = [
+            f"{m.id} ({'private' if m.private else 'public'})"
+            for m in list_models(author=profile.username, token=oauth_token.token)
+        ]
+        return "No models found." if not models else "Models:\n\n" + "\n - ".join(models)
+    except Exception as e:
+        return f"Error listing models: {e}"
+
+
+# --- UTILITIES ---
+
+def get_sdk_version(sdk_choice: str) -> str:
+    pkg = "gradio" if sdk_choice == "gradio" else "streamlit"
+    try:
+        return importlib.metadata.version(pkg)
+    except importlib.metadata.PackageNotFoundError:
+        return "UNKNOWN"
+
+# This utility is no longer needed as parsing is handled in run_codegen
+# def extract_code(text: str) -> str:
+#     blocks = re.findall(r"```(?:\w*\n)?([\s\S]*?)```", text)
+#     return blocks[-1].strip() if blocks else text.strip()
+
+def classify_errors(logs: str) -> str:
+    errs = set()
+    # Convert logs to lower for case-insensitive matching
+    logs_lower = logs.lower()
+    if "syntaxerror" in logs_lower:
+        errs.add("syntax")
+    elif "importerror" in logs_lower or "modulenotfounderror" in logs_lower:
+        errs.add("import")
+    # Catch common error indicators
+    elif "traceback" in logs_lower or "exception" in logs_lower or "error" in logs_lower:
+         errs.add("runtime/generic") # More general error indication
+
+    return ", ".join(errs) or "none"
+
+# --- HF SPACE LOGGING ---
+
+def _get_space_jwt(repo_id: str, token: str) -> str:
+    """Fetches JWT for Space logs using the user's Hf token."""
+    url = f"{constants.ENDPOINT}/api/spaces/{repo_id}/jwt"
+    headers = build_hf_headers(token=token)
+    r = get_session().get(url, headers=headers)
+    hf_raise_for_status(r) # Raises HTTPError for bad responses (e.g. 404 if repo doesn't exist)
+    return r.json()["token"]
+
+def fetch_logs(repo_id: str, level: str, token: str) -> str:
+    """Fetches build or run logs from an HF Space."""
+    try:
+        jwt = _get_space_jwt(repo_id, token)
+        url = f"https://api.hf.space/v1/{repo_id}/logs/{level}"
+        lines = []
+        headers = build_hf_headers(token=jwt)
+        # Use a timeout for the request
+        with get_session().get(url, headers=headers, stream=True, timeout=10) as resp:
+            hf_raise_for_status(resp)
+            # Read lines with a timeout
+            for raw in resp.iter_lines(decode_unicode=True, chunk_size=512):
+                 if raw is None: # handle keep-alive or similar
+                      continue
+                 if raw.startswith("data: "):
+                    try:
+                        ev = json.loads(raw[len("data: "):])
+                        ts, txt = ev.get("timestamp","N/A"), ev.get("data","")
+                        lines.append(f"[{ts}] {txt}")
+                    except json.JSONDecodeError:
+                        lines.append(f"Error decoding log line: {raw}")
+                    except Exception as e:
+                         lines.append(f"Unexpected error processing log line: {raw} - {e}")
+        return "\n".join(lines)
+    except requests.exceptions.Timeout:
+        return f"Error: Timeout fetching {level} logs."
+    except requests.exceptions.RequestException as e:
+        return f"Error fetching {level} logs: {e}"
+    except Exception as e:
+        return f"An unexpected error occurred while fetching logs: {e}"
+
 
-# --- AGENT PROMPTS & FUNCTIONS ---
+def check_iframe(url: str, timeout: int = 10) -> bool:
+    """Checks if the iframe URL is reachable and returns a 200 status code."""
+    try:
+        # Use a HEAD request for efficiency if only status is needed, but GET is safer for
+        # checking if content is served. Let's stick to GET with a timeout.
+        response = requests.get(url, timeout=timeout)
+        return response.status_code == 200
+    except requests.exceptions.RequestException:
+        return False # Any request exception (timeout, connection error, etc.) means it's not accessible
+
+# --- AGENT PROMPTS ---
 
-# Updated/New System Prompts
 SYSTEM_ORCHESTRATOR = {
     "role": "system",
     "content": (
         "You are **Orchestrator Agent**, the project manager. "
         "Your role is to guide the development process from user request to a deployed HF Space application. "
-        "You will delegate tasks to other agents (Architect, Code-Gen, Debug) based on the project state. "
-        "Your inputs will include the current project state (requirements, plan, files, logs, feedback, status). "
-        "Your output should be a decision on the next step/task, potentially summarizing information for the next agent."
-        "Possible next steps could be: 'PLANNING', 'CODING - {task_description}', 'DEBUGGING', 'COMPLETE', 'FAILED'."
-        "Analyze the feedback and logs carefully to decide the next coding task if debugging is needed."
+        "You will analyze the current project state (requirements, plan, files, logs, feedback, status, attempt_count) "
+        "and decide the *single* next step/task for the team. "
+        "Output *only* the name of the next task from the following list: "
+        "'PLANNING', 'CODING - {task_description}', 'PUSHING', 'LOGGING', 'DEBUGGING', 'COMPLETE', 'FAILED'. "
+        "If moving to 'CODING', briefly describe the specific area to focus on (e.g., 'CODING - Initial UI', 'CODING - Adding Data Loading', 'CODING - Fixing Import Errors'). "
+        "Analyze the debug feedback and logs carefully to decide the appropriate coding task description."
+        "If the debug feedback indicates 'All clear', transition to 'COMPLETE'."
+        "If maximum attempts are reached or a critical error occurs, transition to 'FAILED'."
     )
 }
 
@@ -22,9 +135,10 @@ SYSTEM_ARCHITECT = {
     "role": "system",
     "content": (
         "You are **Architect Agent**, the lead planner. "
-        "Given the user requirements, your task is to devise a high-level plan for the application. "
-        "Outline the main features, suggest a file structure (e.g., `app.py`, `utils.py`), and identify key components needed. "
-        "Output the plan clearly, perhaps using bullet points or a simple structure. Do NOT write code."
+        "Given the user requirements and the current project state, your task is to devise or refine the high-level plan for the application. "
+        "Outline the main features, suggest a logical structure, identify potential files (e.g., `app.py`, `utils.py`, `requirements.txt`), and key components needed. "
+        "The target SDK is {sdk_choice}. The main application file should be `{main_app_file}`. "
+        "Output the plan clearly, using bullet points or a simple numbered list. Do NOT write code. Focus only on the plan."
     )
 }
 
@@ -32,12 +146,13 @@ SYSTEM_CODEGEN = {
     "role": "system",
     "content": (
         "You are **Code‑Gen Agent**, a proactive AI developer. "
-        "Your sole responsibility is to author and correct code files based on the plan and assigned task. "
-        "You will receive the full project state, including the plan, existing files, and debug feedback. "
-        "Based on the current task assigned by the Orchestrator, write or modify the necessary code *only* in the specified file(s). "
-        "Output the *full content* of the updated file(s) in markdown code blocks, clearly indicating the filename(s)."
-        "Example output format: \n\n`app.py`\n```python\n# app code here\n```\n\n`utils.py`\n```python\n# utils code here\n```"
-        "Focus *only* on outputting the code blocks for the files you are modifying/creating. Do not add extra commentary outside the code blocks."
+        "Your sole responsibility is to author and correct code files based on the plan and the assigned task. "
+        "You will receive the full project state, including the requirements, plan, existing files, and debug feedback. "
+        "Based on the current task assigned by the Orchestrator ('{current_task}'), write or modify the necessary code *only* in the specified file(s). "
+        "Output the *full content* of the updated file(s) in markdown code blocks, clearly indicating the filename(s) immediately before the code block like this: `filename`\n```<language>\ncode goes here\n```"
+        "If the task involves creating a new file, include it in the output. If modifying an existing file, provide the *complete* modified code for that file."
+        "Ensure the code adheres to the plan and addresses the debug feedback if provided."
+        "Only output the code blocks and their preceding filenames. Do not add extra commentary outside the code blocks."
     )
 }
 
@@ -47,533 +162,621 @@ SYSTEM_DEBUG = {
         "You are **Debug Agent**, a meticulous code reviewer and tester. "
         "You have access to the full project state: requirements, plan, code files, build logs, and run logs. "
         "Your task is to analyze the logs and code in the context of the plan and requirements. "
-        "Identify errors, potential issues, missing features based on the plan, and suggest concrete improvements or fixes. "
-        "If the application appears to be working and meets the plan/requirements (and no significant errors in logs), state 'All clear. Project appears complete.' "
-        "Otherwise, provide actionable feedback for the Code-Gen agent, referencing file names and line numbers where possible. Format feedback clearly."
-        "Example feedback: 'Error in `app.py`: ModuleNotFoundError for 'missing_library'. Add 'missing_library' to `requirements.txt`.\nIssue: The 'plot_data' function in `utils.py` doesn't handle empty input data per the plan.' "
-        "Do NOT write or suggest code directly in your feedback unless it's a very small, inline fix suggestion. Focus on *what* needs fixing/adding and *why*."
+        "Identify errors, potential issues, missing features based on the plan, and suggest concrete improvements or fixes for the Code-Gen agent. "
+        "Pay close attention to the build and run logs for specific errors (SyntaxError, ImportError, runtime errors). "
+        "Also check if the implemented features align with the plan."
+        "If the application appears to be working based on the logs and iframe check, and seems to meet the plan's core requirements, state 'All clear. Project appears complete.' as the *first line* of your feedback."
+        "Otherwise, provide actionable feedback, referencing file names and line numbers where possible. Format feedback clearly."
+        "Example feedback:\n'Error in `app.py`: ModuleNotFoundError for 'missing_library'. Add 'missing_library' to `requirements.txt`.'\n'Issue: The plan required a download button, but it's missing in `app.py`.'\n'Suggestion: Check the loop in `utils.py`, it might cause an infinite loop based on run logs.' "
+        "Do NOT write or suggest large code blocks directly in your feedback. Focus on *what* needs fixing/adding and *why*."
     )
 }
 
-# --- AGENT FUNCTIONS ---
+# --- AGENT RUNNER HELPER ---
 
-def run_agent(client, model_name, system_prompt, user_input, config, chat_history=None):
-    """Helper to run a single agent interaction."""
-    if chat_history is None:
-        chat_history = []
+def run_agent(client, model_name, system_prompt_template, user_input_state, config):
+    """Helper to run a single agent interaction using the project state as input."""
+    # Format the system prompt using state variables if needed
+    try:
+        # Use a specific formatting approach that handles missing keys gracefully if needed
+        # For our current prompts, direct f-string formatting with known keys is fine
+        system_prompt = system_prompt_template["content"].format(**user_input_state)
+    except KeyError as e:
+         print(f"Error formatting system prompt: Missing key {e}. Prompt template: {system_prompt_template['content']}")
+         return f"ERROR: Internal agent error - Missing key {e} for prompt formatting.", None
 
-    messages = chat_history + [{"role": "user", "content": user_input}]
+    # Prepare the message content. The user_input_state is the core context.
+    # We structure it clearly for the agent to read.
+    user_message_content = "Project State:\n" + json.dumps(user_input_state, indent=2)
 
-    # Gemini expects alternating user/model roles starting with user,
-    # but system prompts are outside this. We'll prepend system instructions
-    # to the *first* user message if no history exists, or manage roles manually.
-    # A simpler way for few-turn interactions is to put the system prompt
-    # at the start of the *first* user message, or use a dedicated system role
-    # if the model supports it (Gemini sometimes treats the first entry specially).
-    # Let's stick to prepending system info to the first message for robustness.
-    if not chat_history:
-         user_input = system_prompt["content"] + "\n\n" + user_input
-         messages = [{"role": "user", "content": user_input}] # Start fresh with combined prompt
+    messages = [
+        {"role": "user", "content": system_prompt + "\n\n" + user_message_content}
+    ]
 
-    # For subsequent turns, just append the new user input
-    if chat_history:
-         messages = chat_history + [{"role": "user", "content": user_input}]
+    try:
+        # Use a specific, more capable model if needed for complex reasoning, otherwise Flash is fine
+        # model_to_use = "gemini-1.5-flash-latest" # Example of using a different model
+        model_to_use = model_name # Use the one passed in
 
+        response = client.models.generate_content(
+            model=model_to_use,
+            contents=messages,
+            config=config
+        )
+        response.raise_for_status() # Raise an exception for bad status codes
 
-    # Ensure messages strictly alternate user/model roles for Gemini
-    # This is a common pain point. Let's just send the last few relevant messages + system context
-    # A better approach is to pass the *full* history but carefully structure it.
-    # For simplicity here, let's just pass the system prompt and the immediate user input,
-    # letting the agent rely on the *structured* project_state for context, not chat history.
-    # This shifts complexity from prompt history management to state management.
+        # Some models return parts, concatenate them
+        response_text = "".join([part.text for part in response.candidates[0].content.parts])
 
-    # Revised approach: Pass relevant project_state info *as* the user input
-    # The agent's core prompt (SYSTEM_*) defines its role and what to do with the input state.
-    full_prompt = f"{system_prompt['content']}\n\nProject State:\n{user_input}"
+        print(f"--- Agent Response --- ({model_name})")
+        # print(response_text) # Careful: can be very long
+        print("----------------------")
+
+        return response_text.strip(), messages + [{"role": "model", "content": response_text.strip()}] # Return response and the turn
 
-    try:
-        resp = client.models.generate_content(
-            model=model_name,
-            contents=[{"role": "user", "content": full_prompt}], # Only send the current state input
-            config=config
-        )
-        resp.raise_for_status() # Raise an exception for bad status codes
-        return resp.text.strip(), [{"role": "user", "content": full_prompt}, {"role": "model", "content": resp.text.strip()}] # Return response and the turn
     except Exception as e:
-        print(f"Agent failed: {e}")
-        return f"ERROR: Agent failed - {e}", None # Indicate failure
+        print(f"Agent call failed: {e}")
+        # Attempt to extract error message from response object if possible
+        error_details = str(e)
+        if hasattr(e, 'response') and e.response is not None:
+             try:
+                 error_details = e.response.text # Get response body for more details
+             except:
+                 pass # Ignore if cannot get text
+
+        return f"ERROR: Agent failed - {error_details}", None # Indicate failure
+
+# --- AGENT FUNCTIONS (called by Orchestrator) ---
 
 def run_planner(client, project_state, config):
-    print("Running Planner Agent...")
-    input_state_str = f"Requirements: {project_state['requirements']}"
+    print("Orchestrator: Running Planner Agent...")
+    # Planner needs requirements and basic project info
+    input_state_for_planner = {
+        "requirements": project_state['requirements'],
+        "sdk_choice": project_state['sdk_choice'],
+        "main_app_file": project_state['main_app_file'],
+        # Include other relevant state like existing files if refining plan
+        "files": project_state['files']
+    }
     response_text, _ = run_agent(
         client=client,
-        model_name="gemini-2.5-flash-preview-04-17", # Or a stronger model if needed for planning
-        system_prompt=SYSTEM_ARCHITECT,
-        user_input=input_state_str,
+        model_name="gemini-1.5-flash-latest", # Flash is good for this, maybe Pro if planning is complex
+        system_prompt_template=SYSTEM_ARCHITECT,
+        user_input_state=input_state_for_planner,
         config=config,
-        # chat_history=project_state.get('planner_chat', []) # Optional: keep chat history per agent
     )
-    # project_state['planner_chat'] = planner_chat # Update chat history
 
     if response_text.startswith("ERROR:"):
-        return response_text # Propagate error
+        project_state['status_message'] = response_text
+        return False # Indicate failure
 
     project_state['plan'] = response_text
-    print("Planner Output:", response_state['plan'])
-    return "Plan generated."
+    print("Orchestrator: Planner Output Received.")
+    project_state['status_message'] = "Planning complete."
+    return True
 
 def run_codegen(client, project_state, config):
-    print(f"Running Code-Gen Agent for task: {project_state['current_task']}...")
+    print(f"Orchestrator: Running Code-Gen Agent for task: {project_state['current_task']}...")
     # Code-Gen needs requirements, plan, existing files, and debug feedback
-    input_state_str = (
-        f"Current Task: {project_state['current_task']}\n\n"
-        f"Requirements:\n{project_state['requirements']}\n\n"
-        f"Plan:\n{project_state['plan']}\n\n"
-        f"Existing Files:\n{json.dumps(project_state['files'], indent=2)}\n\n"
-        f"Debug Feedback:\n{project_state['feedback'] or 'None'}\n"
-    )
+    input_state_for_codegen = {
+        "current_task": project_state['current_task'],
+        "requirements": project_state['requirements'],
+        "plan": project_state['plan'],
+        "files": project_state['files'], # Pass current files so it can modify
+        "feedback": project_state['feedback'] or 'None',
+        "sdk_choice": project_state['sdk_choice'],
+        "main_app_file": project_state['main_app_file'] # Ensure it knows the main file convention
+    }
     response_text, _ = run_agent(
         client=client,
-        model_name="gemini-2.5-flash-preview-04-17", # Or a stronger model
-        system_prompt=SYSTEM_CODEGEN,
-        user_input=input_state_str,
+        model_name="gemini-1.5-flash-latest", # Can use Flash, or Pro for more complex coding
+        system_prompt_template=SYSTEM_CODEGEN,
+        user_input_state=input_state_for_codegen,
         config=config,
-        # chat_history=project_state.get('codegen_chat', []) # Optional
     )
-    # project_state['codegen_chat'] = codegen_chat # Update chat history
 
     if response_text.startswith("ERROR:"):
-        return response_text # Propagate error
+        project_state['status_message'] = response_text
+        return False # Indicate failure
 
     # Parse the response text to extract code blocks for potentially multiple files
     files_updated = {}
-    # Simple regex to find blocks like `filename`\n```...\n```
-    # This needs to handle multiple blocks
-    blocks = re.findall(r"(`[^`]*`)\s*```(?:\w*\n)?([\s\S]*?)```", response_text)
+    # Regex to find blocks like `filename`\n```[language]\ncode...\n```
+    # It captures the filename (group 1) and the code content (group 2)
+    # Added handling for optional language tag after the triple backticks
+    blocks = re.findall(r"(`[^`]+`)\s*```(?:\w*\n)?([\s\S]*?)```", response_text)
 
     if not blocks:
-         # If no code blocks are found, maybe the agent just gave commentary or failed
-         # Treat this as an error or a failure to produce code
-         print("Code-Gen Agent did not output any code blocks.")
-         return "ERROR: Code-Gen Agent failed to output code blocks."
+         print("Code-Gen Agent did not output any code blocks in expected format.")
+         project_state['status_message'] = "ERROR: Code-Gen Agent failed to output code blocks in `filename`\\n```code``` format."
+         return False # Indicate failure
 
+    syntax_errors = []
     for filename_match, code_content in blocks:
         filename = filename_match.strip('`').strip()
-        if filename: # Ensure filename is not empty
-            files_updated[filename] = code_content.strip()
+        if not filename:
+             print(f"Warning: Found a code block but filename was empty: {code_content[:50]}...")
+             syntax_errors.append(f"Code block found with empty filename.")
+             continue
+
+        files_updated[filename] = code_content.strip() # Store updated code
+
+        # Quick syntax check for Python files
+        if filename.endswith('.py'):
+            try:
+                compile(code_content, filename, "exec")
+                print(f"Syntax check passed for {filename}")
+            except SyntaxError as e:
+                syntax_errors.append(f"Syntax Error in {filename}: {e}")
+                print(f"Syntax Error in {filename}: {e}")
+            except Exception as e:
+                 syntax_errors.append(f"Unexpected error during syntax check for {filename}: {e}")
+                 print(f"Unexpected error during syntax check for {filename}: {e}")
+
 
     if not files_updated:
-         print("Code-Gen Agent outputted blocks but couldn't parse filenames.")
-         return "ERROR: Code-Gen Agent outputted blocks but couldn't parse filenames."
+         print("Code-Gen Agent outputted blocks but couldn't parse any valid filenames.")
+         project_state['status_message'] = "ERROR: Code-Gen Agent outputted blocks but couldn't parse any valid filenames."
+         return False # Indicate failure
+
+    if syntax_errors:
+        # If syntax errors found, add them to feedback and signal failure for CodeGen step
+        project_state['feedback'] = "Syntax Errors:\n" + "\n".join(syntax_errors) + "\n\n" + project_state['feedback'] # Prepend errors
+        project_state['status_message'] = "ERROR: Code-Gen Agent introduced syntax errors."
+        return False # Indicate failure due to syntax errors
 
 
     project_state['files'].update(files_updated) # Update existing files or add new ones
-    print(f"Code-Gen Agent updated files: {list(files_updated.keys())}")
+    print(f"Orchestrator: Code-Gen Agent updated files: {list(files_updated.keys())}")
 
-    # Quick syntax check on updated files
-    for fn, code in files_updated.items():
-        if fn.endswith('.py'):
-            try:
-                compile(code, fn, "exec")
-                print(f"Syntax check passed for {fn}")
-            except SyntaxError as e:
-                print(f"Syntax Error in {fn}: {e}")
-                return f"ERROR: SyntaxError in {fn}: {e}"
+    # Add the generated/updated code content to the status message for visibility in UI
+    code_summary = "\n".join([f"`{fn}`:\n```python\n{code[:200]}...\n```" for fn, code in files_updated.items()]) # Show snippet
+    project_state['status_message'] = f"Code generated/updated.\n\nFiles Updated:\n{code_summary}"
 
-    return "Code generated/updated."
+    return True # Indicate success
 
 def run_debugger(client, project_state, config):
-    print("Running Debug Agent...")
-    # Debugger needs requirements, plan, files, logs
-    input_state_str = (
-        f"Requirements:\n{project_state['requirements']}\n\n"
-        f"Plan:\n{project_state['plan']}\n\n"
-        f"Current Files:\n{json.dumps(project_state['files'], indent=2)}\n\n"
-        f"Build Logs:\n{project_state['logs'].get('build', 'No build logs.')}\n\n"
-        f"Run Logs:\n{project_state['logs'].get('run', 'No run logs.')}\n"
-        f"Iframe Check: {'Responded OK' if project_state.get('iframe_ok') else 'Not responding or checked yet.'}"
-    )
+    print("Orchestrator: Running Debug Agent...")
+    # Debugger needs requirements, plan, files, logs, and iframe status
+    input_state_for_debugger = {
+        "requirements": project_state['requirements'],
+        "plan": project_state['plan'],
+        "files": project_state['files'],
+        "build_logs": project_state['logs'].get('build', 'No build logs.'),
+        "run_logs": project_state['logs'].get('run', 'No run logs.'),
+        "iframe_status": 'Responding OK' if project_state.get('iframe_ok', False) else 'Not responding or check failed.',
+        "error_types_found": classify_errors(project_state['logs'].get('build', '') + '\n' + project_state['logs'].get('run', ''))
+    }
     response_text, _ = run_agent(
         client=client,
-        model_name="gemini-2.5-flash-preview-04-17", # Flash is good for analysis
-        system_prompt=SYSTEM_DEBUG,
-        user_input=input_state_str,
+        model_name="gemini-1.5-flash-latest", # Flash is good for analysis
+        system_prompt_template=SYSTEM_DEBUG,
+        user_input_state=input_state_for_debugger,
         config=config,
-        # chat_history=project_state.get('debug_chat', []) # Optional
     )
-    # project_state['debug_chat'] = debug_chat # Update chat history
 
     if response_text.startswith("ERROR:"):
-        return response_text # Propagate error
+        project_state['status_message'] = response_text
+        return False # Indicate failure
 
     project_state['feedback'] = response_text
-    print("Debug Agent Feedback:", project_state['feedback'])
-    return "Feedback generated."
-
-def run_orchestrator(client, project_state, config):
-    print("Running Orchestrator Agent...")
-    # The Orchestrator's prompt is implicit in the state it receives and the function logic
-    # We *could* make the Orchestrator an LLM call too, but a state machine is often more reliable
-    # for controlling flow. Let's use Python code as the Orchestrator for now.
-
-    status_message = "Orchestrator deciding next step..."
-
-    # Based on current status and feedback, decide next task
-    if project_state['status'] == 'START':
-        project_state['current_task'] = 'PLANNING'
-        project_state['status'] = 'In Progress'
-        status_message = "Starting project: Planning phase."
-
-    elif project_state['current_task'] == 'PLANNING':
-        # Assume run_planner was just executed
-        if project_state['plan']:
-             project_state['current_task'] = 'CODING - Initial Implementation'
-             status_message = "Planning complete. Moving to initial coding."
-        else:
-             # Planner failed to produce a plan
-             project_state['status'] = 'Failed'
-             status_message = "Planning failed. Could not generate project plan."
-
-    elif project_state['current_task'].startswith('CODING'):
-        # Assume run_codegen was just executed
-        # Next step is always pushing after coding
-        project_state['current_task'] = 'PUSHING'
-        status_message = "Coding phase complete. Preparing to push changes."
-
-    elif project_state['current_task'] == 'PUSHING':
-         # Assume files were just pushed in the main loop
-         # Next step is fetching logs
-         project_state['current_task'] = 'LOGGING'
-         status_message = "Push complete. Fetching logs."
-
-    elif project_state['current_task'] == 'LOGGING':
-        # Assume logs were fetched in the main loop
-        # Next step is debugging
-        project_state['current_task'] = 'DEBUGGING'
-        status_message = "Logs fetched. Analyzing with Debug Agent."
-
-    elif project_state['current_task'] == 'DEBUGGING':
-        # Assume run_debugger was just executed
-        feedback = project_state['feedback']
-        build_logs = project_state['logs'].get('build', '')
-        run_logs = project_state['logs'].get('run', '')
-        iframe_ok = project_state.get('iframe_ok', False)
-
-        # Analyze feedback and logs to decide next step
-        if "All clear. Project appears complete." in feedback or \
-           ("ERROR" not in build_logs.upper() and "ERROR" not in run_logs.upper() and iframe_ok):
-            project_state['status'] = 'Complete'
-            project_state['current_task'] = 'FINISHED'
-            status_message = "Debug Agent reports clear. Project appears complete."
-        elif project_state['attempt_count'] >= 6: # Max attempts reached
-            project_state['status'] = 'Failed'
-            project_state['current_task'] = 'FINISHED'
-            status_message = f"Max attempts ({project_state['attempt_count']}/6) reached. Project failed."
-        else:
-            # Errors or issues found, need more coding/debugging
-            project_state['current_task'] = 'CODING - Addressing Feedback'
-            status_message = "Debug Agent found issues. Returning to Coding phase to address feedback."
-
-    else:
-         # Should not happen, maybe a cleanup or failure state
-         project_state['status'] = 'Failed'
-         status_message = f"Orchestrator in unknown state: {project_state['current_task']}"
-
-
-    project_state['status_message'] = status_message
-    print(f"Orchestrator Status: {project_state['status']} - Task: {project_state['current_task']}")
+    print("Orchestrator: Debug Agent Feedback Received.")
+    project_state['status_message'] = "Debug feedback generated."
+    return True
 
+# --- MAIN ORCHESTRATION LOGIC ---
 
-# --- MAIN HANDLER REVISION ---
-
-def handle_user_message(
-    history, # Keep history for chatbot display, but agents use project_state
-    sdk_choice: str,
-    gemini_api_key: str,
-    grounding_enabled: bool,
-    temperature: float,
-    max_output_tokens: int,
-    profile: gr.OAuthProfile | None,
-    oauth_token: gr.OAuthToken | None
-):
-    if not profile or not oauth_token:
-        return history + [{"role":"assistant","content":"⚠️ Please log in first."}], "", "", "<p>No Space yet.</p>", "Waiting for prompt..."
+# We use Python functions to represent the state transitions and agent calls
+# This is more deterministic than using an LLM purely for orchestration.
 
-    client = genai.Client(api_key=gemini_api_key)
-    repo_id = f"{profile.username}/{profile.username}-auto-space"
-    iframe_url = f"https://huggingface.co/spaces/{repo_id}"
-    sdk_version = get_sdk_version(sdk_choice)
-    code_fn = "app.py" if sdk_choice == "gradio" else "streamlit_app.py" # Still need a primary file name
+def orchestrate_development(client, project_state, config, oauth_token_token):
+    """Manages the overall development workflow."""
 
-    # Initialize or load project state
-    # For this implementation, we'll initialize fresh for each new user message
-    # A more complex app might persist state or load from HF Space
-    project_state = {
-        'requirements': history[-1]['content'] if history else "", # Get the latest user message
-        'plan': '',
-        'files': {}, # Use a dict to store multiple file contents {filename: code}
-        'logs': {'build': '', 'run': ''},
-        'feedback': '',
-        'current_task': 'START',
-        'status': 'In Progress',
-        'status_message': 'Initializing...',
-        'attempt_count': 0,
-        'sdk_choice': sdk_choice,
-        'sdk_version': sdk_version,
-        'repo_id': repo_id,
-        'iframe_url': iframe_url,
-        'main_app_file': code_fn, # Keep track of the main app file name convention
-        # Optional: Per-agent chat histories for more complex interactions
-        # 'planner_chat': [],
-        # 'codegen_chat': [],
-        # 'debug_chat': [],
-    }
-
-    # Add initial user message to history for chatbot display
-    if history and history[-1].get("role") == "user":
-        # This assumes the Gradio chatbot already added the user message
-        pass # History should already have the user input
-    else:
-        # Handle cases where history might be empty or last turn wasn't user
-        # (e.g. first message, or if previous turn failed)
-        # This depends on how Gradio handles the initial history input
-        # For standard Gradio chatbot, history will be like [{"role": "user", "content": "..."}]
-        pass
-
-
-    cfg = GenerateContentConfig(
-        tools=[Tool(google_search=GoogleSearch())] if grounding_enabled else [],
-        response_modalities=["TEXT"],
-        temperature=temperature,
-        max_output_tokens=max_output_tokens, # Ensure this is high enough for code
-    )
-
-    # Main Orchestration Loop
     while project_state['status'] == 'In Progress' and project_state['attempt_count'] < 7:
         print(f"\n--- Attempt {project_state['attempt_count'] + 1} ---")
         print(f"Current Task: {project_state['current_task']}")
+
         current_task = project_state['current_task']
-        feedback_message = "" # Message to add to chatbot history about step outcome
+        step_successful = True # Flag to track if the current step completed without error
+        # Add current task to history for UI visibility
+        project_state['chat_history'].append({"role": "assistant", "content": f"➡️ Task: {current_task}"})
+        project_state['status_message'] = f"Executing: {current_task}..."
+
 
         if current_task == 'START':
-             run_orchestrator(client, project_state, cfg) # Move to PLANNING
+            # Initial state, move to planning
+            project_state['current_task'] = 'PLANNING'
+            project_state['status_message'] = "Starting project: Initializing."
+
 
         elif current_task == 'PLANNING':
-             result = run_planner(client, project_state, cfg)
-             feedback_message = f"**Plan:**\n\n{project_state['plan']}"
-             if result.startswith("ERROR:"):
-                 project_state['status'] = 'Failed'
-                 project_state['status_message'] = result
-             run_orchestrator(client, project_state, cfg) # Decide next based on plan result
+            step_successful = run_planner(client, project_state, config)
+            if step_successful:
+                 project_state['current_task'] = 'CODING - Initial Implementation' # Move to coding after planning
+                 # Add plan to chat history for user
+                 project_state['chat_history'].append({"role": "assistant", "content": f"**Plan:**\n{project_state['plan']}"})
+            else:
+                 project_state['current_task'] = 'FAILED' # Planning failed
 
-        elif current_task.startswith('CODING'):
-             # If this is the very first coding task, ensure the main app file exists
-             # even if the planner didn't explicitly list it.
-             if project_state['main_app_file'] not in project_state['files'] and project_state['attempt_count'] == 0:
-                 project_state['files'][project_state['main_app_file']] = f"# Initial {sdk_choice} app file\n" # Start with something
 
-             result = run_codegen(client, project_state, cfg)
-             if result.startswith("ERROR:"):
-                 project_state['status'] = 'Failed'
-                 project_state['status_message'] = result
-             else:
-                 # Show the code that was just generated/updated in the chatbot history
-                 code_output = "\n".join([f"`{fn}`\n```python\n{code}\n```" for fn, code in project_state['files'].items()])
-                 feedback_message = f"**Generated/Updated Code:**\n\n{code_output}"
+        elif current_task.startswith('CODING'):
+            # Ensure the main app file exists before coding if it's the first coding step
+            if project_state['attempt_count'] == 0 and project_state['current_task'] == 'CODING - Initial Implementation':
+                 if project_state['main_app_file'] not in project_state['files']:
+                     project_state['files'][project_state['main_app_file']] = f"# Initial {project_state['sdk_choice']} app file\n" # Start with a basic stub
+                     if project_state['sdk_choice'] == 'gradio':
+                          project_state['files'][project_state['main_app_file']] += "import gradio as gr\n\n# Define a simple interface\n# For example: gr.Interface(...).launch()\n"
+                     elif project_state['sdk_choice'] == 'streamlit':
+                           project_state['files'][project_state['main_app_file']] += "import streamlit as st\n\n# Your Streamlit app starts here\n# For example: st.write('Hello, world!')\n"
+
+
+            step_successful = run_codegen(client, project_state, config)
+
+            if step_successful:
+                 project_state['current_task'] = 'PUSHING' # Always push after attempting to code
+            else:
+                 # Code-gen failed (syntax error, parsing issue, etc.)
+                 # We'll try debugging/coding again in the next attempt loop iteration if attempts allow
+                 print("Code-Gen step failed. Incrementing attempt count.")
+                 project_state['attempt_count'] += 1 # Count Code-Gen failure as an attempt
+                 # Add the error message from run_codegen to chat history
+                 project_state['chat_history'].append({"role": "assistant", "content": project_state['status_message']})
+                 project_state['current_task'] = 'DEBUGGING' # Go to debugging to analyze the failure
 
-             run_orchestrator(client, project_state, cfg) # Move to PUSHING or Failed
 
         elif current_task == 'PUSHING':
-            # Save and push all files in project_state['files']
             try:
                 # Create/update repo first
-                create_repo(repo_id=repo_id, token=oauth_token.token,
-                            exist_ok=True, repo_type="space", space_sdk=sdk_choice)
+                create_repo(repo_id=project_state['repo_id'], token=oauth_token_token,
+                            exist_ok=True, repo_type="space", space_sdk=project_state['sdk_choice'])
 
-                # Write and upload all files
+                # Write and upload all files in project_state['files']
                 for fn, content in project_state['files'].items():
-                     # Write to a temporary file or in the current directory
+                     # Ensure directories exist for files like utils/data.py
+                     os.makedirs(os.path.dirname(fn), exist_ok=True)
                      with open(fn, "w") as f:
                          f.write(content)
                      upload_file(
                          path_or_fileobj=fn, path_in_repo=fn,
-                         repo_id=repo_id, token=oauth_token.token,
+                         repo_id=project_state['repo_id'], token=oauth_token_token,
                          repo_type="space"
                      )
 
-                # Ensure requirements.txt and README are also pushed, or handled by CodeGen
-                # For simplicity, let's add boilerplate here if not generated by CodeGen
+                # Ensure requirements.txt is handled - CodeGen *could* generate it,
+                # but adding boilerplate here guarantees it exists if not generated.
                 if 'requirements.txt' not in project_state['files']:
-                     req_content = "pandas\n" + ("streamlit\n" if sdk_choice=="streamlit" else "gradio\n")
+                     req_content = "pandas\n" + ("streamlit\n" if project_state['sdk_choice']=="streamlit" else "gradio\n") + "google-generativeai\nhuggingface-hub\n"
                      with open("requirements.txt", "w") as f:
                          f.write(req_content)
                      upload_file(path_or_fileobj="requirements.txt", path_in_repo="requirements.txt",
-                                 repo_id=repo_id, token=oauth_token.token, repo_type="space")
+                                 repo_id=project_state['repo_id'], token=oauth_token_token, repo_type="space")
+                     project_state['files']['requirements.txt'] = req_content # Add to state for completeness
 
+                # Ensure README.md is handled
                 if 'README.md' not in project_state['files']:
                     readme_content = f"""---
-title: {profile.username}'s Auto Space
+title: {project_state['repo_id']}
 emoji: 🐢
-sdk: {sdk_choice}
-sdk_version: {sdk_version}
+sdk: {project_state['sdk_choice']}
+sdk_version: {project_state['sdk_version']}
 app_file: {project_state['main_app_file']}
 pinned: false
 ---
+# {project_state['repo_id']}
+
+This is an auto-generated HF Space.
+
+**Requirements:** {project_state['requirements']}
+
+**Plan:**
+{project_state['plan']}
 """
                     with open("README.md", "w") as f:
                         f.write(readme_content)
                     upload_file(path_or_fileobj="README.md", path_in_repo="README.md",
-                                 repo_id=repo_id, token=oauth_token.token, repo_type="space")
+                                 repo_id=project_state['repo_id'], token=oauth_token_token, repo_type="space")
+                    project_state['files']['README.md'] = readme_content # Add to state
 
 
-                print(f"Pushed {len(project_state['files'])} files to {repo_id}")
-                feedback_message = f"Pushed code to HF Space **{repo_id}**."
-                run_orchestrator(client, project_state, cfg) # Move to LOGGING
+                print(f"Pushed {len(project_state['files'])} files to {project_state['repo_id']}")
+                project_state['status_message'] = f"Pushed code to HF Space **{project_state['repo_id']}**. Waiting for build..."
+                project_state['chat_history'].append({"role": "assistant", "content": project_state['status_message']})
+                project_state['current_task'] = 'LOGGING' # Move to fetching logs
 
             except Exception as e:
-                project_state['status'] = 'Failed'
+                step_successful = False
+                project_state['status'] = 'Failed' # Pushing is critical, fail if it fails
                 project_state['status_message'] = f"ERROR: Failed to push to HF Space: {e}"
-                feedback_message = project_state['status_message']
-                print(feedback_message)
+                project_state['chat_history'].append({"role": "assistant", "content": project_state['status_message']})
+                print(project_state['status_message'])
+                project_state['current_task'] = 'FINISHED' # End process
 
 
         elif current_task == 'LOGGING':
-            # Fetch logs and check iframe status
-            try:
-                project_state['logs']['build'] = fetch_logs(repo_id, "build")
-                project_state['logs']['run'] = fetch_logs(repo_id, "run")
-                project_state['iframe_ok'] = check_iframe(iframe_url)
-                feedback_message = "**Logs Fetched.** Analyzing..."
-                print("Logs fetched.")
-                run_orchestrator(client, project_state, cfg) # Move to DEBUGGING
-            except Exception as e:
-                 project_state['status'] = 'Failed'
-                 project_state['status_message'] = f"ERROR: Failed to fetch logs or check iframe: {e}"
-                 feedback_message = project_state['status_message']
-                 print(feedback_message)
+            # Wait a moment for build to start
+            time.sleep(5) # Initial wait
+            wait_time = 5
+            max_log_wait = 60 # Maximum total time to wait for logs in this step
+            elapsed_log_wait = 0
+            logs_fetched = False
+
+            while elapsed_log_wait < max_log_wait:
+                 try:
+                     build_logs = fetch_logs(project_state['repo_id'], "build", oauth_token_token)
+                     run_logs = fetch_logs(project_state['repo_id'], "run", oauth_token_token)
+                     project_state['logs']['build'] = build_logs
+                     project_state['logs']['run'] = run_logs
+                     project_state['iframe_ok'] = check_iframe(project_state['iframe_url'])
+
+                     logs_fetched = True
+                     print("Logs fetched. Checking iframe.")
+
+                     # Simple check: if build logs exist AND (contain success or timeout reached)
+                     # and run logs exist, we can proceed to debugging.
+                     # More advanced: check for specific build success/failure patterns.
+                     if logs_fetched and (
+                          "Building success" in build_logs or elapsed_log_wait >= max_log_wait - wait_time # Assume built if timeout reached
+                          or "Build complete" in build_logs # Common build tool output
+                         ) and (run_logs or project_state['iframe_ok']): # Need some sign of runtime or iframe
+                            break # Exit the log fetching wait loop
+                     elif "ERROR" in build_logs.upper() or "FATAL" in build_logs.upper():
+                          print("Build errors detected, proceeding to debugging.")
+                          break # Proceed to debugging to analyze errors
+                     else:
+                         print(f"Logs incomplete or no clear status yet. Waiting {wait_time}s...")
+                         time.sleep(wait_time)
+                         elapsed_log_wait += wait_time
+                         wait_time = min(wait_time * 1.5, 15) # Increase wait time
+
+
+                 except Exception as e:
+                      print(f"Error during log fetching or iframe check: {e}. Will retry.")
+                      time.sleep(wait_time)
+                      elapsed_log_wait += wait_time
+                      wait_time = min(wait_time * 1.5, 15)
+
+
+            if logs_fetched:
+                project_state['status_message'] = "Logs fetched and iframe checked."
+                project_state['chat_history'].append({"role": "assistant", "content": project_state['status_message']})
+                project_state['current_task'] = 'DEBUGGING' # Move to debugging to analyze logs
+            else:
+                 step_successful = False
+                 project_state['status'] = 'Failed' # Failed to fetch logs within timeout
+                 project_state['status_message'] = "ERROR: Failed to fetch logs or iframe within timeout."
+                 project_state['chat_history'].append({"role": "assistant", "content": project_state['status_message']})
+                 project_state['current_task'] = 'FINISHED' # End process
 
 
         elif current_task == 'DEBUGGING':
-            result = run_debugger(client, project_state, cfg)
-            feedback_message = f"**Debug Feedback:**\n\n{project_state['feedback']}"
-            if result.startswith("ERROR:"):
+            step_successful = run_debugger(client, project_state, config)
+
+            # Add debug feedback to chat history
+            project_state['chat_history'].append({"role": "assistant", "content": f"**Debug Feedback:**\n{project_state['feedback']}"})
+
+
+            if step_successful:
+                 # Analyze feedback to decide next step
+                 feedback = project_state['feedback']
+                 iframe_ok = project_state.get('iframe_ok', False)
+                 error_types = classify_errors(project_state['logs'].get('build', '') + '\n' + project_state['logs'].get('run', ''))
+
+                 print(f"Debug Analysis - Feedback: {feedback[:100]}... | Iframe OK: {iframe_ok} | Errors: {error_types}")
+
+
+                 if "All clear. Project appears complete." in feedback or \
+                    (iframe_ok and error_types == "none" and "ERROR" not in feedback.upper()):
+                     # Debugger says it's clear OR iframe is working AND no errors reported in logs AND no ERROR in feedback
+                     project_state['status'] = 'Complete'
+                     project_state['current_task'] = 'FINISHED'
+                     project_state['status_message'] = "Debug Agent reports clear. Project appears complete."
+                 elif project_state['attempt_count'] >= 6: # Max attempts reached AFTER debugging
+                     project_state['status'] = 'Failed'
+                     project_state['current_task'] = 'FINISHED'
+                     project_state['status_message'] = f"Max attempts ({project_state['attempt_count']+1}/7) reached after debugging. Project failed."
+                 else:
+                     # Errors or issues found, need more coding/debugging
+                     project_state['current_task'] = 'CODING - Addressing Feedback'
+                     project_state['status_message'] = "Debug Agent found issues. Returning to Coding phase to address feedback."
+                     project_state['attempt_count'] += 1 # Increment attempt count before looping back to coding
+                     backoff_wait = min(project_state['attempt_count'] * 5, 30) # Backoff before next coding attempt
+                     print(f"Waiting {backoff_wait} seconds before next coding attempt...")
+                     time.sleep(backoff_wait)
+
+            else:
+                 # Debugger failed (e.g. API error)
                  project_state['status'] = 'Failed'
-                 project_state['status_message'] = result
-            run_orchestrator(client, project_state, cfg) # Move to Decision/CODING/COMPLETE/FAILED
+                 project_state['current_task'] = 'FINISHED'
+                 # status_message already set by run_debugger
 
-        elif current_task == 'FINISHED':
-            # Exit loop state
-            pass
 
-        else:
-            # Unknown task
-            project_state['status'] = 'Failed'
-            project_state['status_message'] = f"ERROR: Orchestrator entered an unknown task state: {current_task}"
-            feedback_message = project_state['status_message']
-            print(feedback_message)
+        elif current_task == 'FINISHED':
+            # Exit the main loop
+            break
 
+        # --- End of Task Handling ---
 
-        # Add the outcome of the last step to the chatbot history
-        if feedback_message:
-             history.append({"role": "assistant", "content": feedback_message})
-             # Provide status updates as part of the assistant messages or via the status_message output
-             history.append({"role": "assistant", "content": f"Project Status: {project_state['status_message']}"})
+        # If a step failed and didn't already set status to FAILED, mark the Orchestrator logic as failed
+        if not step_successful and project_state['status'] == 'In Progress':
+             project_state['status'] = 'Failed'
+             project_state['status_message'] = project_state.get('status_message', 'An unexpected error occurred during orchestration step.')
+             project_state['current_task'] = 'FINISHED' # End process
 
+    # --- End of Orchestration Loop ---
 
-        # Increment attempt count after a cycle of Coding/Pushing/Logging/Debugging
-        # Only increment if we went through the core loop steps, not planning or initial start
-        if current_task in ['DEBUGGING', 'PLANNING']: # End of a major cycle
-             project_state['attempt_count'] += 1
-             # Add a small backoff before the next attempt loop starts
-             if project_state['status'] == 'In Progress':
-                 backoff = min(project_state['attempt_count'] * 5, 30) # Linear or exponential backoff
-                 print(f"Waiting for {backoff} seconds before next attempt...")
-                 time.sleep(backoff)
+    # Final status message if loop exited without explicit FINISHED state
+    if project_state['status'] == 'In Progress':
+         project_state['status'] = 'Failed'
+         project_state['status_message'] = project_state.get('status_message', 'Orchestration loop exited unexpectedly.')
 
 
-    # --- End of Orchestration Loop ---
+    # Add final outcome message to history
+    project_state['chat_history'].append({"role": "assistant", "content": f"**Project Outcome:** {project_state['status']} - {project_state['status_message']}"})
+    if project_state['status'] == 'Complete':
+        project_state['chat_history'].append({"role": "assistant", "content": "✅ Application deployed successfully (likely)! Check the preview above."})
+    else:
+         project_state['chat_history'].append({"role": "assistant", "content": "❌ Project failed to complete. Review logs and feedback for details."})
 
-    # Prepare final UI outputs
-    final_messages = history[:] # Use the accumulated history
-    final_build_logs = project_state['logs'].get('build', '')
-    final_run_logs = project_state['logs'].get('run', '')
 
-    # Final iframe check and HTML
-    final_iframe_html = (
-        f'<iframe src="{iframe_url}" width="100%" height="500px"></iframe>'
-        + ("" if project_state.get('iframe_ok') else "<p style='color:red;'>⚠️ iframe not responding or checked.</p>")
+    # Return final state for UI update
+    return (
+        project_state['chat_history'],
+        project_state['logs'].get('build', 'No build logs.'),
+        project_state['logs'].get('run', 'No run logs.'),
+        (f'<iframe src="{project_state["iframe_url"]}" width="100%" height="500px"></iframe>'
+         + ("" if project_state.get('iframe_ok') else "<p style='color:red;'>⚠️ iframe not responding or checked.</p>")),
+        project_state['status_message'] # Return the final status message
     )
 
-    # Add a final status message to the chatbot history
-    final_messages.append({"role": "assistant", "content": f"**Project Outcome:** {project_state['status']} - {project_state['status_message']}"})
-    if project_state['status'] == 'Complete':
-        final_messages.append({"role": "assistant", "content": "✅ Application deployed successfully (likely)! Check the preview above."})
-    else:
-         final_messages.append({"role": "assistant", "content": "❌ Project failed to complete. Review logs and feedback for details."})
 
-    return final_messages, final_build_logs, final_run_logs, final_iframe_html, project_state['status_message']
+# --- MAIN HANDLER (Called by Gradio) ---
 
+def handle_user_message(
+    history, # This is the list of messages in the Gradio Chatbot
+    sdk_choice: str,
+    gemini_api_key: str,
+    grounding_enabled: bool,
+    temperature: float,
+    max_output_tokens: int,
+    profile: gr.OAuthProfile | None,
+    oauth_token: gr.OAuthToken | None # We need the token object
+):
+    if not profile or not oauth_token:
+        # Append error message to history for display
+        error_history = history + [{"role":"assistant","content":"⚠️ Please log in first."}]
+        return error_history, "", "", "<p>Please log in.</p>", "Login required."
 
-# --- SIMPLE UI WITH HIGHER MAX TOKENS & STATUS DISPLAY —
+    if not gemini_api_key:
+         error_history = history + [{"role":"assistant","content":"⚠️ Please provide your Gemini API Key."}]
+         return error_history, "", "", "<p>Please provide API Key.</p>", "API Key required."
+
+
+    client = genai.Client(api_key=gemini_api_key)
+    repo_id = f"{profile.username}/{profile.username}-auto-space"
+    iframe_url = f"https://huggingface.co/spaces/{repo_id}"
+    sdk_version = get_sdk_version(sdk_choice)
+    code_fn = "app.py" if sdk_choice == "gradio" else "streamlit_app.py" # Standard main file name convention
+
+    # Get the user's latest prompt from the history
+    user_prompt = history[-1]['content'] if history and history[-1].get("role") == "user" else "No prompt provided."
+    if user_prompt == "No prompt provided." and len(history) > 0:
+         # Handle cases where history might be unusual, try finding last user message
+         for msg in reversed(history):
+             if msg.get("role") == "user":
+                 user_prompt = msg["content"]
+                 break
+
+
+    # Initialize project state for this development session
+    # History will be updated throughout and returned at the end
+    project_state = {
+        'requirements': user_prompt,
+        'plan': '',
+        'files': {}, # Use a dict to store multiple file contents {filename: code}
+        'logs': {'build': '', 'run': ''},
+        'feedback': '',
+        'current_task': 'START',
+        'status': 'In Progress',
+        'status_message': 'Initializing...',
+        'attempt_count': 0,
+        'sdk_choice': sdk_choice,
+        'sdk_version': sdk_version,
+        'repo_id': repo_id,
+        'iframe_url': iframe_url,
+        'main_app_file': code_fn,
+        'chat_history': history[:] # Use the passed-in history to build upon
+    }
+
+    cfg = GenerateContentConfig(
+        tools=[Tool(google_search=GoogleSearch())] if grounding_enabled else [],
+        response_modalities=["TEXT"],
+        temperature=temperature,
+        max_output_tokens=int(max_output_tokens), # Ensure integer
+    )
+
+    # Start the orchestration process
+    final_history, final_build_logs, final_run_logs, final_iframe_html, final_status_message = orchestrate_development(
+        client, project_state, cfg, oauth_token.token
+    )
+
+    # Return the final state for the UI
+    return (
+        final_history,
+        final_build_logs,
+        final_run_logs,
+        final_iframe_html,
+        final_status_message
+    )
+
+# --- SIMPLE UI WITH HIGHER MAX TOKENS & STATUS DISPLAY ---
 
 with gr.Blocks(title="HF Space Auto‑Builder (Team AI)") as demo:
-    gr.Markdown("## 🐢 HF Space Auto‑Builder (Team AI)\n1) Sign in  2) Enter prompt  3) Watch progress, logs & preview")
+    gr.Markdown("## 🐢 HF Space Auto‑Builder (Team AI)\nUse AI agents to build and deploy a simple Gradio or Streamlit app on a Hugging Face Space.")
+    gr.Markdown("1) Log in with Hugging Face. 2) Enter your Gemini API Key. 3) Provide app requirements. 4) Click 'Start Development Team' and watch the process.")
 
-    login_btn = gr.LoginButton(variant="huggingface", size="lg")
-    status_md = gr.Markdown("*Not logged in.*")
-    models_md = gr.Markdown()
-    demo.load(show_profile,     inputs=None, outputs=status_md)
-    demo.load(list_private_models, inputs=None, outputs=models_md)
-    login_btn.click(show_profile,     inputs=None, outputs=status_md)
-    login_btn.click(list_private_models, inputs=None, outputs=models_md)
 
     with gr.Row():
         with gr.Column(scale=1):
-            sdk_choice = gr.Radio(["gradio","streamlit"], value="gradio", label="SDK")
-            api_key    = gr.Textbox(label="Gemini API Key", type="password")
-            grounding  = gr.Checkbox(label="Enable grounding", value=False)
-            temp       = gr.Slider(0,1,value=0.2, label="Temperature")
-            max_tokens = gr.Number(value=4096, label="Max output tokens (recommend >= 4096)", minimum=512) # Increased recommendation
+            login_btn = gr.LoginButton(variant="huggingface", size="lg")
+            status_md = gr.Markdown("*Not logged in.*")
+            models_md = gr.Markdown()
+            # Load profile/models on startup and login button click
+            demo.load(show_profile, inputs=None, outputs=status_md, api_name="load_profile")
+            demo.load(list_private_models, inputs=[login_btn, login_btn], outputs=models_md, api_name="load_models")
+            login_btn.select(show_profile, inputs=[login_btn], outputs=status_md, api_name="login_profile") # Use select for login_btn
+            login_btn.select(list_private_models, inputs=[login_btn, login_btn], outputs=models_md, api_name="login_models") # Pass login_btn state twice to get profile and token
+
+
+            gr.Markdown("---")
+
+            sdk_choice = gr.Radio(["gradio","streamlit"], value="gradio", label="SDK", info="Choose the framework for your app.")
+            api_key    = gr.Textbox(label="Gemini API Key", type="password", info="Get one from Google AI Studio.")
+            grounding  = gr.Checkbox(label="Enable Google Search (Grounding)", value=False, info="Allow agents to use Google Search.")
+            temp       = gr.Slider(0,1,value=0.2, label="Temperature", info="Creativity of agents. Lower is more focused.")
+            max_tokens = gr.Number(value=4096, label="Max Output Tokens", minimum=1000, info="Max length of agent responses (code, feedback, etc.). Recommend 4096+.")
 
         with gr.Column(scale=2):
             project_status_md = gr.Markdown("Waiting for prompt...")
-            chatbot   = gr.Chatbot(type="messages", label="Team Communication & Status")
-            user_in   = gr.Textbox(placeholder="Your prompt for the application...", label="Application Requirements", lines=2)
-            send_btn  = gr.Button("Start Development Team")
+            chatbot   = gr.Chatbot(type="messages", label="Team Communication & Status", show_copy_button=True)
+            user_in   = gr.Textbox(placeholder="Describe the application you want to build...", label="Application Requirements", lines=3)
+            send_btn  = gr.Button("🚀 Start Development Team")
 
     # Separate accordions for logs and preview
     with gr.Accordion("Logs", open=False):
-        build_box = gr.Textbox(label="Build logs", lines=5, interactive=False)
-        run_box   = gr.Textbox(label="Run logs",   lines=5, interactive=False)
-        refresh_btn = gr.Button("Refresh Logs Only") # Add refresh button for logs
+        build_box = gr.Textbox(label="Build logs", lines=10, interactive=False, max_lines=20)
+        run_box   = gr.Textbox(label="Run logs",   lines=10, interactive=False, max_lines=20)
+        # Need login state for refresh button
+        refresh_btn = gr.Button("🔄 Refresh Logs Only")
 
     with gr.Accordion("App Preview", open=True):
-        preview   = gr.HTML("<p>No Space yet.</p>")
-
+        preview   = gr.HTML("<p>App preview will load here when available.</p>")
 
     # Update the button click handler
     # It will now return the updated chatbot history, logs, preview, and the project status
     send_btn.click(
         fn=handle_user_message,
-        inputs=[chatbot, sdk_choice, api_key, grounding, temp, max_tokens, login_btn, login_btn],
+        inputs=[chatbot, sdk_choice, api_key, grounding, temp, max_tokens, login_btn, login_btn], # Pass login_btn state twice
         outputs=[chatbot, build_box, run_box, preview, project_status_md]
     )
     user_in.submit(
-         fn=handle_user_message,
-        inputs=[chatbot, sdk_choice, api_key, grounding, temp, max_tokens, login_btn, login_btn],
+        fn=handle_user_message,
+        inputs=[chatbot, sdk_choice, api_key, grounding, temp, max_tokens, login_btn, login_btn], # Pass login_btn state twice
         outputs=[chatbot, build_box, run_box, preview, project_status_md]
     )
 
     # Handler for refreshing logs manually
+    # It needs the repo_id (derived from profile) and the auth token
     refresh_btn.click(
         fn=lambda profile, token: (
-            fetch_logs(f"{profile.username}/{profile.username}-auto-space", "build"),
-            fetch_logs(f"{profile.username}/{profile.username}-auto-space", "run")
+            fetch_logs(f"{profile.username}/{profile.username}-auto-space", "build", token.token) if profile and token else "Login required to fetch logs.",
+            fetch_logs(f"{profile.username}/{profile.username}-auto-space", "run", token.token) if profile and token else "Login required to fetch logs."
         ),
-        inputs=[login_btn, login_btn], # Pass profile and token from login button state
+        inputs=[login_btn, login_btn], # Pass login_btn state twice to get profile and token
         outputs=[build_box, run_box]
     )
 
 
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+    # Clean up files created during the process when the app stops (optional, good for Spaces)
+    # demo.on_event("close", lambda: [os.remove(f) for f in os.listdir() if os.path.isfile(f) and (f.endswith(".py") or f.endswith(".txt") or f.endswith(".md"))]) # Be careful with this in production
+
+demo.launch(server_name="0.0.0.0", server_port=7860)