Update app.py

app.py

@@ -24,6 +24,10 @@ def list_private_models(
     if not profile or not oauth_token:
         return "Please log in to see your models."
     try:
+        # Check if the token is valid before listing models
+        if not hasattr(oauth_token, 'token') or not oauth_token.token:
+             return "Invalid or missing access token."
+
         models = [
             f"{m.id} ({'private' if m.private else 'public'})"
             for m in list_models(author=profile.username, token=oauth_token.token)
@@ -42,11 +46,6 @@ def get_sdk_version(sdk_choice: str) -> str:
     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
@@ -73,6 +72,8 @@ def _get_space_jwt(repo_id: str, token: str) -> str:
 
 def fetch_logs(repo_id: str, level: str, token: str) -> str:
     """Fetches build or run logs from an HF Space."""
+    if not token:
+         return "Login required to fetch logs."
     try:
         jwt = _get_space_jwt(repo_id, token)
         url = f"https://api.hf.space/v1/{repo_id}/logs/{level}"
@@ -208,7 +209,7 @@ def run_agent(client, model_name, system_prompt_template, user_input_state, conf
         # Some models return parts, concatenate them
         response_text = "".join([part.text for part in response.candidates[0].content.parts])
 
-        print(f"--- Agent Response --- ({model_name})")
+        print(f"--- Agent Response --- ({model_to_use})")
         # print(response_text) # Careful: can be very long
         print("----------------------")
 
@@ -289,6 +290,9 @@ def run_codegen(client, project_state, config):
     if not 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."
+         # Add the agent's raw response to chat history for debugging if no blocks found
+         project_state['chat_history'].append({"role": "assistant", "content": f"Code-Gen Agent raw response (no code blocks detected):\n\n{response_text}"})
+
          return False # Indicate failure
 
     syntax_errors = []
@@ -321,8 +325,11 @@ def run_codegen(client, project_state, config):
 
     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."
+        project_state['feedback'] = "Syntax Errors Found:\n" + "\n".join(syntax_errors) + "\n\n" + project_state['feedback'] # Prepend errors
+        project_state['status_message'] = "ERROR: Code-Gen Agent introduced syntax errors. Debugging needed."
+        # Add syntax errors to chat history for user visibility
+        project_state['chat_history'].append({"role": "assistant", "content": project_state['status_message']})
+        project_state['chat_history'].append({"role": "assistant", "content": "Details:\n" + "\n".join(syntax_errors)})
         return False # Indicate failure due to syntax errors
 
 
@@ -330,8 +337,10 @@ def run_codegen(client, project_state, config):
     print(f"Orchestrator: Code-Gen Agent updated files: {list(files_updated.keys())}")
 
     # 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}"
+    code_summary = "\n".join([f"`{fn}`:\n```python\n{code[:500]}{'...' if len(code) > 500 else ''}\n```" for fn, code in files_updated.items()]) # Show snippet
+    project_state['chat_history'].append({"role": "assistant", "content": f"**Code Generated/Updated:**\n\n{code_summary}"})
+    project_state['status_message'] = f"Code generated/updated: {list(files_updated.keys())}"
+
 
     return True # Indicate success
 
@@ -362,6 +371,7 @@ def run_debugger(client, project_state, config):
     project_state['feedback'] = response_text
     print("Orchestrator: Debug Agent Feedback Received.")
     project_state['status_message'] = "Debug feedback generated."
+    # Add debug feedback to chat history (already done in orchestrate_development, but could be here)
     return True
 
 # --- MAIN ORCHESTRATION LOGIC ---
@@ -372,42 +382,63 @@ def run_debugger(client, project_state, config):
 def orchestrate_development(client, project_state, config, oauth_token_token):
     """Manages the overall development workflow."""
 
+    # Initial step transition
+    if project_state['current_task'] == 'START':
+        project_state['current_task'] = 'PLANNING'
+        project_state['status_message'] = "Starting project: Initializing and moving to Planning."
+
+
     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']
-        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}..."
 
+        # Add current task to history for UI visibility (if not already added by previous step)
+        if not project_state['chat_history'] or project_state['chat_history'][-1].get('content', '').strip() != f"➡️ Task: {current_task}".strip():
+             project_state['chat_history'].append({"role": "assistant", "content": f"➡️ Task: {current_task}"})
 
-        if current_task == 'START':
-            # Initial state, move to planning
-            project_state['current_task'] = 'PLANNING'
-            project_state['status_message'] = "Starting project: Initializing."
-
+        step_successful = True # Flag to track if the current step completed without error
 
-        elif current_task == 'PLANNING':
+        if current_task == 'PLANNING':
             step_successful = run_planner(client, project_state, config)
             if step_successful:
+                 # Add plan to chat history for user (done inside run_planner now?) - moved to run_planner
                  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'):
-            # 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':
+            # Ensure the main app file exists before coding if it's the first coding step (or first after syntax error)
+            if project_state['attempt_count'] == 0 or 'Syntax Errors' in project_state.get('feedback', ''):
                  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"
+                     # Also ensure requirements.txt and README are stubbed if they don't exist
+                     if 'requirements.txt' not in project_state['files']:
+                          project_state['files']['requirements.txt'] = "pandas\n" + ("streamlit\n" if project_state['sdk_choice']=="streamlit" else "gradio\n") + "google-generativeai\nhuggingface-hub\n"
+                     if 'README.md' not in project_state['files']:
+                         readme_content = f"""---
+title: {project_state['repo_id']}
+emoji: 🐢
+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']}
+"""
+                         project_state['files']['README.md'] = readme_content
 
 
             step_successful = run_codegen(client, project_state, config)
@@ -416,11 +447,10 @@ def orchestrate_development(client, project_state, config, oauth_token_token):
                  project_state['current_task'] = 'PUSHING' # Always push after attempting to code
             else:
                  # Code-gen failed (syntax error, parsing issue, etc.)
+                 # The failure is handled within run_codegen by setting status_message and feedback
                  # 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']})
+                 print("Code-Gen step failed.")
+                 # attempt_count is incremented AFTER debugging phase analyses results
                  project_state['current_task'] = 'DEBUGGING' # Go to debugging to analyze the failure
 
 
@@ -442,42 +472,6 @@ def orchestrate_development(client, project_state, config, oauth_token_token):
                          repo_type="space"
                      )
 
-                # 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 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=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: {project_state['repo_id']}
-emoji: 🐢
-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=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 {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']})
@@ -496,9 +490,14 @@ This is an auto-generated HF Space.
             # 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
+            max_log_wait = 90 # Maximum total time to wait for logs in this step (increased slightly)
             elapsed_log_wait = 0
             logs_fetched = False
+            iframe_checked = False
+
+            project_state['status_message'] = "Fetching logs and checking iframe..."
+            project_state['chat_history'].append({"role": "assistant", "content": project_state['status_message']})
+
 
             while elapsed_log_wait < max_log_wait:
                  try:
@@ -506,44 +505,50 @@ This is an auto-generated HF Space.
                      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
+
+                     # Only check iframe once logs indicate something might be running, or after a delay
+                     if elapsed_log_wait > 10: # Don't check immediately
+                         project_state['iframe_ok'] = check_iframe(project_state['iframe_url'])
+                         iframe_checked = True
+                     else:
+                         project_state['iframe_ok'] = False # Assume not ready yet
+
+
+                     print(f"Log/Iframe check at {elapsed_log_wait}s. Build logs len: {len(build_logs)}, Run logs len: {len(run_logs)}, Iframe OK: {project_state['iframe_ok']}")
+
+                     # Conditions to proceed to debugging:
+                     # 1. Iframe is OK (app is running and accessible)
+                     # 2. Build logs show errors (need debugging ASAP)
+                     # 3. Max wait time is almost reached (proceed with whatever logs we have)
+                     # 4. Build logs exist and indicate *some* progress (e.g., contain "Building" or sufficient length)
+                     if project_state['iframe_ok'] or \
+                        "ERROR" in build_logs.upper() or "FATAL" in build_logs.upper() or \
+                        elapsed_log_wait >= max_log_wait - wait_time or \
+                        ("Building" in build_logs or len(build_logs) > 100) and logs_fetched: # Heuristic for build progress
                             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...")
+                         print(f"Logs or iframe not ready. Waiting {wait_time}s...")
                          time.sleep(wait_time)
                          elapsed_log_wait += wait_time
-                         wait_time = min(wait_time * 1.5, 15) # Increase wait time
+                         wait_time = min(wait_time * 1.5, 20) # Increase wait time, cap at 20s
 
 
                  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)
+                      wait_time = min(wait_time * 1.5, 20)
 
 
-            if logs_fetched:
-                project_state['status_message'] = "Logs fetched and iframe checked."
+            if logs_fetched or iframe_checked: # Proceed if we got logs OR checked the iframe
+                project_state['status_message'] = "Logs fetched and iframe checked (or timeout reached)."
                 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['status'] = 'Failed' # Failed to fetch logs/check iframe within timeout after retries
+                 project_state['status_message'] = "ERROR: Failed to fetch logs or check iframe within timeout after multiple retries."
                  project_state['chat_history'].append({"role": "assistant", "content": project_state['status_message']})
                  project_state['current_task'] = 'FINISHED' # End process
 
@@ -559,8 +564,17 @@ This is an auto-generated HF Space.
                  # Analyze feedback to decide next step
                  feedback = project_state['feedback']
                  iframe_ok = project_state.get('iframe_ok', False)
+                 # Re-check logs just before making the decision for freshest info if possible
+                 # This might add latency, skipping for now, rely on logs fetched in LOGGING step
+                 # build_logs = fetch_logs(project_state['repo_id'], "build", oauth_token_token)
+                 # run_logs = fetch_logs(project_state['repo_id'], "run", oauth_token_token)
+                 # error_types = classify_errors(build_logs + '\n' + run_logs)
+                 # project_state['logs']['build'] = build_logs # Update state with freshest logs
+                 # project_state['logs']['run'] = run_logs
+
                  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}")
 
 
@@ -570,7 +584,7 @@ This is an auto-generated HF Space.
                      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
+                 elif project_state['attempt_count'] >= 6: # Max attempts reached AFTER debugging analysis
                      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."
@@ -578,7 +592,7 @@ This is an auto-generated HF Space.
                      # 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
+                     project_state['attempt_count'] += 1 # Increment attempt count AFTER a debug cycle points 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)
@@ -592,30 +606,49 @@ This is an auto-generated HF Space.
 
         elif current_task == 'FINISHED':
             # Exit the main loop
-            break
+            pass # Loop condition handles exit
 
-        # --- End of Task Handling ---
+        else:
+            # Unknown task
+            step_successful = False
+            project_state['status'] = 'Failed'
+            project_state['status_message'] = f"ERROR: Orchestrator entered an unknown task state: {current_task}"
+            project_state['chat_history'].append({"role": "assistant", "content": project_state['status_message']})
+            print(project_state['status_message'])
+            project_state['current_task'] = 'FINISHED' # End process
 
-        # If a step failed and didn't already set status to FAILED, mark the Orchestrator logic as failed
+
+        # If a step failed and didn't explicitly set status to FAILED (like PUSHING),
+        # the orchestrator logic above should handle transition to FAILED or DEBUGGING.
+        # This check acts as a safeguard.
         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['status_message'] = project_state.get('status_message', f'An error occurred during task: {current_task}')
              project_state['current_task'] = 'FINISHED' # End process
 
+
     # --- End of Orchestration Loop ---
 
     # Final status message if loop exited without explicit FINISHED state
     if project_state['status'] == 'In Progress':
+         # This shouldn't happen with the current logic, but good practice
          project_state['status'] = 'Failed'
          project_state['status_message'] = project_state.get('status_message', 'Orchestration loop exited unexpectedly.')
 
 
-    # Add final outcome message to history
-    project_state['chat_history'].append({"role": "assistant", "content": f"**Project Outcome:** {project_state['status']} - {project_state['status_message']}"})
+    # Add final outcome message to history if not already the last message
+    final_outcome_message = f"**Project Outcome:** {project_state['status']} - {project_state['status_message']}"
+    if not project_state['chat_history'] or project_state['chat_history'][-1].get('content', '').strip() != final_outcome_message.strip():
+        project_state['chat_history'].append({"role": "assistant", "content": final_outcome_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."})
+        completion_message = "✅ Application deployed successfully (likely)! Check the preview above."
+        if not project_state['chat_history'] or project_state['chat_history'][-1].get('content', '').strip() != completion_message.strip():
+             project_state['chat_history'].append({"role": "assistant", "content": completion_message})
+    elif project_state['status'] == 'Failed':
+         failure_message = "❌ Project failed to complete. Review logs and feedback for details."
+         if not project_state['chat_history'] or project_state['chat_history'][-1].get('content', '').strip() != failure_message.strip():
+             project_state['chat_history'].append({"role": "assistant", "content": failure_message})
 
 
     # Return final state for UI update
@@ -624,7 +657,7 @@ This is an auto-generated HF Space.
         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>")),
+         + ("" if project_state.get('iframe_ok') else "<p style='color:red;'>⚠️ iframe not responding or check failed.</p>")),
         project_state['status_message'] # Return the final status message
     )
 
@@ -641,9 +674,21 @@ def handle_user_message(
     profile: gr.OAuthProfile | None,
     oauth_token: gr.OAuthToken | None # We need the token object
 ):
-    if not profile or not oauth_token:
+    # Append the user's message to the history immediately for display
+    # Gradio's Chatbot usually handles this automatically, but explicitly ensures it's present
+    # before we start the orchestration.
+    # Check if the last message is the current user input to avoid duplication on submit/click
+    if not history or history[-1].get("content") != user_in.value: # Assuming user_in is accessible or passed
+         # This logic is tricky with Gradio's submit button.
+         # The default Chatbot behavior on submit/click is to add the user message.
+         # So, history *should* already contain the new user message.
+         # Let's rely on Gradio's standard behavior.
+         pass
+
+    if not profile or not oauth_token or not oauth_token.token:
         # Append error message to history for display
-        error_history = history + [{"role":"assistant","content":"⚠️ Please log in first."}]
+        error_history = history + [{"role":"assistant","content":"⚠️ Please log in first via the Hugging Face button."}]
+        # Return current state, logs etc. + the new history
         return error_history, "", "", "<p>Please log in.</p>", "Login required."
 
     if not gemini_api_key:
@@ -658,13 +703,13 @@ def handle_user_message(
     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
+    # Assume the last message in history is the user's new prompt because Gradio added it
     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
+
+    if user_prompt == "No prompt provided." or user_prompt.strip() == "":
+        # Handle empty prompt case
+        empty_prompt_history = history + [{"role":"assistant","content":"Please enter requirements for the application."}]
+        return empty_prompt_history, "", "", "<p>Enter requirements.</p>", "Waiting for prompt."
 
 
     # Initialize project state for this development session
@@ -675,7 +720,7 @@ def handle_user_message(
         'files': {}, # Use a dict to store multiple file contents {filename: code}
         'logs': {'build': '', 'run': ''},
         'feedback': '',
-        'current_task': 'START',
+        'current_task': 'START', # Start the orchestration state machine
         'status': 'In Progress',
         'status_message': 'Initializing...',
         'attempt_count': 0,
@@ -696,7 +741,7 @@ def handle_user_message(
 
     # 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
+        client, project_state, cfg, oauth_token.token # Pass the token string
     )
 
     # Return the final state for the UI
@@ -717,14 +762,29 @@ with gr.Blocks(title="HF Space Auto‑Builder (Team AI)") as demo:
 
     with gr.Row():
         with gr.Column(scale=1):
+            # --- LOGIN BUTTON / PROFILE & MODEL LISTING (FIXED) ---
             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
+
+            # On app load, show “not logged in” and list public models (or none)
             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
+            demo.load(list_private_models, inputs=None, outputs=models_md, api_name="load_models") # Inputs=None allows Gradio to inject profile/token (will be None initially)
+
+            # When the user actually logs in:
+            login_btn.click(
+                fn=show_profile,
+                inputs=None,               # Gradio will inject OAuthProfile
+                outputs=status_md,
+                api_name="login_profile"
+            )
+            login_btn.click(
+                fn=list_private_models,
+                inputs=None,               # Gradio will inject (OAuthProfile, OAuthToken)
+                outputs=models_md,
+                api_name="login_models"
+            )
+            # --- END FIX ---
 
 
             gr.Markdown("---")
@@ -745,38 +805,41 @@ with gr.Blocks(title="HF Space Auto‑Builder (Team AI)") as demo:
     with gr.Accordion("Logs", open=False):
         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
+        # Need login state for refresh button, inputs=None will inject profile/token
         refresh_btn = gr.Button("🔄 Refresh Logs Only")
 
     with gr.Accordion("App Preview", open=True):
         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
+    # Inputs=None allows Gradio to inject profile/token from login_btn
     send_btn.click(
         fn=handle_user_message,
-        inputs=[chatbot, sdk_choice, api_key, grounding, temp, max_tokens, login_btn, login_btn], # Pass login_btn state twice
+        inputs=[chatbot, sdk_choice, api_key, grounding, temp, max_tokens, login_btn, login_btn], # Keep login_btn twice for profile and token objects
         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], # Pass login_btn state twice
+        inputs=[chatbot, sdk_choice, api_key, grounding, temp, max_tokens, login_btn, login_btn], # Keep login_btn 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
+    # Inputs=None allows Gradio to inject profile/token from login_btn
     refresh_btn.click(
         fn=lambda profile, token: (
-            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."
+            fetch_logs(f"{profile.username}/{profile.username}-auto-space", "build", token.token) if profile and token and token.token else "Login required to fetch logs.",
+            fetch_logs(f"{profile.username}/{profile.username}-auto-space", "run", token.token) if profile and token and token.token else "Login required to fetch logs."
         ),
-        inputs=[login_btn, login_btn], # Pass login_btn state twice to get profile and token
+        inputs=[login_btn, login_btn], # Keep login_btn twice for profile and token objects
         outputs=[build_box, run_box]
     )
 
 
     # Clean up files created during the process when the app stops (optional, good for Spaces)
+    # Consider adding more specific cleanup if needed
     # 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)