Update app.py

app.py

@@ -517,7 +517,8 @@ def orchestrate_development(client, project_state: dict, config, oauth_token_tok
          project_state['files'] = {}
          project_state['iframe_ok'] = False
          project_state['log_error_state'] = 'None'
-         project_state['chat_history'].append({"role": "assistant", "content": "Project initialized. Starting development team."})
+         # Initial message handled in handle_user_message now
+         # project_state['chat_history'].append({"role": "assistant", "content": "Project initialized. Starting development team."})
          return # Exit after initialization, next call will run the first task
 
     if project_state['status'] != 'In Progress':
@@ -535,7 +536,7 @@ def orchestrate_development(client, project_state: dict, config, oauth_token_tok
     if not project_state['chat_history'] or project_state['chat_history'][-1].get('content', '').strip() != task_message.strip():
          project_state['chat_history'].append({"role": "assistant", "content": task_message})
 
-    step_successful = True # Assume success until a step explicitly fails
+    # step_successful is managed within each task block
 
     if current_task == 'START':
          # This state should only be hit once for initialization, handled above.
@@ -599,7 +600,7 @@ This is an auto-generated HF Space.
              project_state['current_task'] = 'PUSHING'
              # Clear previous logs and feedback before pushing a new version
              project_state['logs'] = {'build': '', 'run': ''}
-             project_state['feedback'] = ''
+             project_state['feedback'] = '' # Clear feedback from previous debug cycle
              project_state['iframe_ok'] = False
              project_state['log_error_state'] = 'None' # Reset log error state
              # Do NOT reset attempt count here, it's for total cycles
@@ -615,7 +616,6 @@ This is an auto-generated HF Space.
              project_state['status_message'] = "ERROR: Cannot push without Hugging Face token."
              project_state['chat_history'].append({"role": "assistant", "content": project_state['status_message']})
              project_state['current_task'] = 'FINISHED'
-             step_successful = False
              print(project_state['status_message'])
              return # Exit early on critical error
 
@@ -642,6 +642,7 @@ This is an auto-generated HF Space.
 
             # Use a temporary directory to avoid cluttering the working directory
             import tempfile
+            step_successful = True # Assume success initially for pushing
             with tempfile.TemporaryDirectory() as tmpdir:
                  temp_file_paths = []
                  for fn, content in files_to_push.items():
@@ -654,39 +655,45 @@ This is an auto-generated HF Space.
                           temp_file_paths.append(full_path)
                       except Exception as e:
                           print(f"Error writing temporary file {fn}: {e}")
-                          # Decide how to handle: skip file? Fail push? For now, skip.
+                          # Decide how to handle: skip file? Fail push? For now, mark push as failed.
+                          step_successful = False # Mark push as failed if any file fails to write
                           project_state['feedback'] = project_state.get('feedback', '') + f"\n\nError writing file {fn} for push: {e}"
+                          project_state['status_message'] = f"ERROR: Failed to write {fn} locally for push. Details in feedback."
 
 
-                 if not temp_file_paths:
+                 if not temp_file_paths and len(files_to_push) > 0: # If files were supposed to be pushed but none were written
                       print("No valid files prepared for push.")
-                      project_state['status_message'] = "ERROR: No files were prepared for push after Code-Gen. Project Failed?"
+                      project_state['status_message'] = "ERROR: No files were prepared for push after Code-Gen. Project Failed."
                       project_state['chat_history'].append({"role": "assistant", "content": project_state['status_message']})
                       project_state['status'] = 'Failed'
                       project_state['current_task'] = 'FINISHED'
                       step_successful = False
-                      return
-
-
-                 # Upload files one by one
-                 for fn in files_to_push.keys(): # Iterate over keys to get path_in_repo
-                     filepath = os.path.join(tmpdir, fn)
-                     if filepath in temp_file_paths: # Only try to upload if temp file was successfully written
-                         try:
-                             upload_file(
-                                 path_or_fileobj=filepath, path_in_repo=fn,
-                                 repo_id=repo_id, token=oauth_token_token,
-                                 repo_type="space"
-                             )
-                             print(f"Uploaded: {fn}")
-                         except Exception as e:
-                              print(f"Error uploading file {fn}: {e}")
-                              step_successful = False # Mark push as failed if any upload fails
-                              project_state['feedback'] = project_state.get('feedback', '') + f"\n\nError uploading {fn}: {e}"
-                              project_state['status_message'] = f"ERROR: Failed to upload {fn}. Details in feedback."
-                              # Continue trying other files, but mark overall step as failed
-                     else:
-                         print(f"Skipping upload for {fn} due to previous write error.")
+                      # No need to return here, let the rest of the block handle the failure state
+
+
+                 # Upload files one by one IF step_successful is still True after writing
+                 if step_successful:
+                      for fn in files_to_push.keys(): # Iterate over keys to get path_in_repo
+                          filepath = os.path.join(tmpdir, fn)
+                          # Only try to upload if temp file was successfully written and overall push isn't already failed
+                          if filepath in temp_file_paths and step_successful:
+                              try:
+                                  upload_file(
+                                      path_or_fileobj=filepath, path_in_repo=fn,
+                                      repo_id=repo_id, token=oauth_token_token,
+                                      repo_type="space"
+                                  )
+                                  print(f"Uploaded: {fn}")
+                              except Exception as e:
+                                   print(f"Error uploading file {fn}: {e}")
+                                   step_successful = False # Mark push as failed if any upload fails
+                                   project_state['feedback'] = project_state.get('feedback', '') + f"\n\nError uploading {fn}: {e}"
+                                   project_state['status_message'] = f"ERROR: Failed to upload {fn}. Details in feedback."
+                                   # Continue trying other files, but mark overall step as failed
+                          elif step_successful: # If step_successful is True but file wasn't written
+                                print(f"Skipping upload for {fn} as it failed to write locally.")
+                                # This case is already covered by the write loop failing step_successful
+
 
             if step_successful:
                  project_state['status_message'] = f"Pushed code to HF Space **{repo_id}**. Build triggered. Waiting for build and logs..."
@@ -696,11 +703,13 @@ This is an auto-generated HF Space.
                  # Logs, iframe_ok, log_error_state were reset before push, will be updated by poller
 
             else:
-                 # If any upload failed, the push step failed
+                 # If any write or upload failed, the push step failed
                  project_state['status'] = 'Failed'
                  project_state['current_task'] = 'FINISHED'
-                 project_state['status_message'] = project_state.get('status_message', f"ERROR: One or more files failed to upload to HF Space {repo_id}. See feedback.")
-                 project_state['chat_history'].append({"role": "assistant", "content": project_state['status_message']})
+                 project_state['status_message'] = project_state.get('status_message', f"ERROR: One or more files failed to process or upload to HF Space {repo_id}. See feedback.")
+                 # Ensure the error message is added to chat history if not already there from writing/uploading
+                 if not project_state['chat_history'] or project_state['chat_history'][-1].get('content', '').strip() != project_state['status_message'].strip():
+                     project_state['chat_history'].append({"role": "assistant", "content": project_state['status_message']})
                  print(project_state['status_message'])
 
 
@@ -747,16 +756,16 @@ This is an auto-generated HF Space.
              project_state['status_message'] = "Iframe check passed. Analyzing build/run logs."
 
         # Check for explicit log fetching errors
-        elif log_error_state != 'None':
+        elif log_error_state not in ['None', 'No Logs Yet', 'Iframe Not Ready']: # Any error state that is NOT just 'waiting'
              print(f"Orchestrator: Detected log fetching error state: {log_error_state}. Moving to Debugging.")
              ready_to_debug = True
              project_state['status_message'] = f"Log fetching encountered an error ({log_error_state}). Moving to Debugging."
 
         # Check for significant logs that indicate build/run started/failed
         # Avoid triggering debug if it's just the 'No logs found yet' message
-        elif not (build_logs.startswith("No build logs") or run_logs.startswith("No run logs")):
+        elif not (build_logs.strip() == "No build logs found yet." and run_logs.strip() == "No run logs found yet."):
             # Check for actual log content or error indicators within logs
-            if len(build_logs) > 50 or len(run_logs) > 10 or "ERROR" in build_logs.upper() or "FATAL" in build_logs.upper() or "ERROR" in run_logs.upper() or "FATAL" in run_logs.upper():
+            if len(build_logs.strip()) > 50 or len(run_logs.strip()) > 10 or "ERROR" in build_logs.upper() or "FATAL" in build_logs.upper() or "ERROR" in run_logs.upper() or "FATAL" in run_logs.upper():
                  print("Orchestrator: Detected significant logs. Moving to Debugging.")
                  ready_to_debug = True
                  project_state['status_message'] = "Significant logs detected. Analyzing build/run logs."
@@ -765,20 +774,17 @@ This is an auto-generated HF Space.
         # Timeout check - Use attempt count as a global timeout for the entire process,
         # but also use it to limit time spent *just* in LOGGING if no progress is detected.
         if project_state['attempt_count'] >= 6: # Use attempt count as a global timeout trigger
-            print("Orchestrator: Max attempts reached in LOGGING. Forcing move to Debugging.")
+            print("Orchestrator: Max attempts reached in LOGGING/overall. Forcing move to Debugging.")
             ready_to_debug = True # Force debug to get final feedback before failing
 
 
         if ready_to_debug:
              project_state['current_task'] = 'DEBUGGING'
              # Add a small delay before debugging starts, giving UI time to update logs
-             time.sleep(2) # Add a slight pause
-
+             # time.sleep(2) # Moved implicit wait to poller tick
         else:
              # If not ready to debug, stay in LOGGING. The orchestrator loop will naturally pause between calls.
-             # Increment attempt count only when staying in LOGGING without progress *after* the initial state.
-             # The attempt count is more for the overall process, but we'll use it to track cycles spent here.
-             # The attempt count is incremented globally at the end of the loop if not finished/failed.
+             # Attempt count is incremented globally at the end of the loop if not finished/failed.
              pass # Stay in LOGGING, loop will repeat
 
 
@@ -817,9 +823,9 @@ This is an auto-generated HF Space.
                   is_failed = True
                   project_state['status_message'] = f"Max attempts ({project_state['attempt_count']+1}/7) reached. Project failed."
                   print(f"Debug Analysis: Failed due to max attempts ({project_state['attempt_count']+1}).")
-             elif log_error_state in ['Auth Failed', 'Auth Failed JWT', 'Auth Failed STREAM', 'Request Error', 'Unexpected Error']: # Critical log fetching errors
+             elif log_error_state in ['Auth Failed', 'Auth Failed JWT', 'Auth Failed STREAM', 'Request Error', 'Unexpected Error', 'Poller Unexpected Error', 'HTTP Error']: # Critical log fetching errors
                   is_failed = True
-                  project_state['status_message'] = f"Project failed due to critical log fetching error: {log_error_state}"
+                  project_state['status_message'] = f"Project failed due to critical log or iframe fetching error: {log_error_state}"
                   print(f"Debug Analysis: Failed due to critical log error: {log_error_state}.")
              elif ("ERROR" in feedback.upper() or error_types != "none") and project_state['attempt_count'] >= 4:
                   # If debugger finds errors or classified errors exist after several attempts (e.g., attempts 5, 6, 7)
@@ -857,7 +863,7 @@ This is an auto-generated HF Space.
 
     else:
         # Should not happen - unknown state
-        step_successful = False # Force failure path
+        # step_successful = False # Force failure path - not needed as status is set
         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']})
@@ -873,6 +879,7 @@ This is an auto-generated HF Space.
         # time.sleep(1) # Adding a 1-second pause between steps
 
     # Final outcome message when status changes to Complete or Failed
+    # Only add if it's the final state and the message hasn't been added
     if project_state['status'] != 'In Progress' and not any(msg.get('content', '').strip().startswith("**Project Outcome:") for msg in project_state['chat_history']):
          final_outcome_message = f"**Project Outcome:** {project_state['status']} - {project_state.get('status_message', 'No specific outcome message.')}"
          project_state['chat_history'].append({"role": "assistant", "content": final_outcome_message})
@@ -887,7 +894,7 @@ This is an auto-generated HF Space.
                   project_state['chat_history'].append({"role": "assistant", "content": failure_message})
 
 
-    # orchestration_development function doesn't return UI outputs directly anymore.
+    # orchestrate_development function doesn't return UI outputs directly anymore.
     # It updates the project_state dictionary, which is a gr.State,
     # and the UI outputs are updated by the poller (_update_logs_and_preview)
     # or by the main handle_user_message function after calling orchestrate_development.
@@ -916,159 +923,147 @@ def _update_logs_and_preview(profile_token_state: tuple[gr.OAuthProfile | None,
     rid    = f"{profile.username}/{clean}"
     url    = f"https://huggingface.co/spaces/{profile.username}/{clean}"
 
-    # Ensure project_state is initialized if poller runs before the first handle_user_message
-    # This shouldn't happen in the intended flow, but defensively check.
+    # Initialize project_state if it's the first call or if space name changed while not mid-project
     if current_project_state is None:
+         print(f"Poller: Initializing state for space {rid}")
          current_project_state = {
              'repo_id': rid, # Set repo_id based on current UI inputs
              'iframe_url': url, # Set iframe_url based on current UI inputs
-             'logs': {'build': 'Initializing...', 'run': 'Initializing...'},
+             'logs': {'build': 'Polling...', 'run': 'Polling...'},
              'iframe_ok': False,
-             'log_error_state': 'Initializing...',
-             'chat_history': [],
+             'log_error_state': 'Polling...', # Initial polling state
+             'chat_history': [], # Start with empty history for a new space
              'status': 'Not Started',
-             'status_message': 'Enter requirements to start.',
+             'status_message': f'Polling logs for {rid}... Enter requirements to start project.',
              'attempt_count': 0,
              'requirements': '', 'plan': '', 'files': {}, 'feedback': '', 'current_task': None, 'sdk_choice': '', 'sdk_version': '', 'main_app_file': ''
          }
+         # Add an initial message if history is empty
+         if not current_project_state['chat_history']:
+              current_project_state['chat_history'].append({"role": "assistant", "content": current_project_state['status_message']})
+
     elif current_project_state.get('repo_id') != rid:
          # If the UI Space name changes while a project is in progress for a *different* space,
          # reset the project state to focus on the new space name.
-         # Alternatively, could just fetch logs for the new space without resetting the whole state.
-         # Resetting seems safer to avoid confusion.
          print(f"Poller detected space name change from {current_project_state.get('repo_id')} to {rid}. Resetting state.")
          current_project_state = {
              'repo_id': rid,
              'iframe_url': url,
              'logs': {'build': 'Space name changed, fetching logs...', 'run': 'Space name changed, fetching logs...'},
              'iframe_ok': False,
-             'log_error_state': 'None',
+             'log_error_state': 'None', # Reset error state
              'chat_history': [{"role": "assistant", "content": f"Space name changed to **{rid}**. Project state reset. Enter requirements to start development on this Space."}],
              'status': 'Not Started',
              'status_message': 'Space name changed. Enter requirements to start.',
-             'attempt_count': 0,
+             'attempt_count': 0, # Reset attempts for a new project
              'requirements': '', 'plan': '', 'files': {}, 'feedback': '', 'current_task': None, 'sdk_choice': '', 'sdk_version': '', 'main_app_file': ''
          }
     else:
-         # If space name matches the one in state, update status message to indicate polling is active
-         if current_project_state.get('status') == 'In Progress' and current_project_state.get('current_task') == 'LOGGING':
-             # Orchestrator updates status message in LOGGING, poller shouldn't overwrite it here.
-             pass # Let orchestrator control message during LOGGING task
-         elif current_project_state.get('status') in ['Complete', 'Failed']:
-              # Don't change status message if project is finished
-              pass
-         else:
-             # For other states (Not Started, In Progress but not LOGGING task), indicate polling
-             current_status_message = f"Polling logs for {rid}... Current status: {current_project_state.get('status_message', '...')}"
-             if not current_project_state['chat_history'] or current_project_state['chat_history'][-1].get('content', '').strip() != current_status_message.strip():
-                 # Avoid flooding chat history, maybe just update project_status_md?
-                 # Let's update state message, UI will show it.
-                 current_project_state['status_message'] = current_status_message
+         # If space name matches the one in state, potentially update polling status message
+         if current_project_state.get('status') in ['Not Started']:
+              current_status_message = f"Polling logs for {rid}... Current status: {current_project_state.get('status_message', '...')}"
+              # Avoid flooding chat history with this, just update the status message field
+              current_project_state['status_message'] = current_status_message
+              # Ensure there's at least an initial message in history if needed
+              if not current_project_state['chat_history']:
+                   current_project_state['chat_history'].append({"role": "assistant", "content": current_project_state['status_message']})
+         # If status is 'In Progress' and task is 'LOGGING', the orchestrator manages status_message.
+         # If status is 'Complete' or 'Failed', the final status_message is set.
+         # Otherwise, keep the current status_message.
 
 
     # Fetch logs and check iframe status using the revised functions
-    build_logs_content = current_project_state['logs'].get('build', 'Fetching...')
-    run_logs_content = current_project_state['logs'].get('run', 'Fetching...')
-    iframe_status_ok = current_project_state.get('iframe_ok', False)
-    current_log_error_state = current_project_state.get('log_error_state', 'None')
+    # Always attempt to fetch logs and check iframe if a repo_id is set in the state.
+    # The fetch_logs function handles errors like 401/404/timeout.
+    if current_project_state.get('repo_id'):
+         print(f"Poller fetching logs for {rid} (Status: {current_project_state.get('status')}, Task: {current_project_state.get('current_task')})...")
 
+         build_logs_content = fetch_logs(rid, "build")
+         run_logs_content = fetch_logs(rid, "run") # Fetch run logs regardless of build log status
+         iframe_status_ok = check_iframe(url)
+
+         # Update log_error_state based on fetch results
+         # Priority to critical errors > iframe issues > waiting
+         new_log_error_state = 'None'
+
+         if build_logs_content.startswith("ERROR_AUTH_FAILED_JWT:") or run_logs_content.startswith("ERROR_AUTH_FAILED_JWT:"):
+             new_log_error_state = 'Auth Failed JWT'
+         elif build_logs_content.startswith("ERROR_AUTH_FAILED_STREAM:") or run_logs_content.startswith("ERROR_AUTH_FAILED_STREAM:"):
+              new_log_error_state = 'Auth Failed STREAM'
+         elif build_logs_content.startswith("ERROR_NOT_FOUND:") or run_logs_content.startswith("ERROR_NOT_FOUND:"):
+             new_log_error_state = 'Not Found'
+         elif build_logs_content.startswith("ERROR_TIMEOUT:") or run_logs_content.startswith("ERROR_TIMEOUT:"):
+              new_log_error_state = 'TIMEOUT'
+         elif build_logs_content.startswith("ERROR_REQUEST:") or run_logs_content.startswith("ERROR_REQUEST:"):
+              new_log_error_state = 'Request Error'
+         elif build_logs_content.startswith("ERROR_HTTP_") or run_logs_content.startswith("ERROR_HTTP_"):
+              new_log_error_state = 'HTTP Error'
+         elif build_logs_content.startswith("ERROR_UNEXPECTED:") or run_logs_content.startswith("ERROR_UNEXPECTED:"):
+              new_log_error_state = 'Poller Unexpected Error'
+         elif not iframe_status_ok:
+              # If iframe is not ok, and no explicit log error, assume it's not ready
+              new_log_error_state = 'Iframe Not Ready'
+         elif build_logs_content.strip() == "No build logs found yet." and run_logs_content.strip() == "No run logs found yet.":
+             # If no logs and iframe ok, maybe it's a very simple app? Or a very fast build?
+             # Let's not mark it as 'No Logs Yet' if iframe is OK.
+             pass # Keep None
+         elif (build_logs_content.strip() == "No build logs found yet." or run_logs_content.strip() == "No run logs found yet.") and iframe_status_ok:
+             # Partial logs, but iframe is OK
+             pass # Keep None
+         elif (build_logs_content.strip() != "No build logs found yet." or run_logs_content.strip() != "No run logs found yet.") and not iframe_status_ok:
+             # Logs are appearing but iframe is not OK
+             new_log_error_state = 'Iframe Not Ready' # Re-assert iframe not ready if logs appear but it's still down
 
-    # Only fetch logs/check iframe if the project is 'In Progress' and the task is 'LOGGING',
-    # OR if the status is 'Not Started' but a space name is entered (to show logs immediately).
-    # OR if status is Complete/Failed, to allow checking final state.
-    is_orchestration_polling = current_project_state.get('status') == 'In Progress' and current_project_state.get('current_task') == 'LOGGING'
-    is_initial_or_finished_poll = current_project_state.get('status') in ['Not Started', 'Complete', 'Failed']
 
-    if is_orchestration_polling or is_initial_or_finished_poll:
-         print(f"Poller fetching logs for {rid} (Status: {current_project_state.get('status')}, Task: {current_project_state.get('current_task')})...")
-         try:
-             build_logs_content = fetch_logs(rid, "build")
-             # Check for specific error markers returned by fetch_logs
-             if build_logs_content.startswith("ERROR_"):
-                 current_log_error_state = build_logs_content.split(':')[0].replace('ERROR_', '') # Extract error type
-             elif current_log_error_state == 'None' and build_logs_content.strip() == "No build logs found yet.":
-                 current_log_error_state = 'No Logs Yet' # Indicate waiting
-             elif current_log_error_state == 'No Logs Yet' and build_logs_content.strip() != "No build logs found yet.":
-                  current_log_error_state = 'None' # Logs started appearing
-
-         except Exception as e: # Should be caught by fetch_logs internal try/except
-             build_logs_content = f"Poller Unexpected error fetching build logs: {e}"
-             current_log_error_state = 'Poller Unexpected Error'
-             print(build_logs_content)
-
-
-         # Only fetch run logs if a critical error wasn't found for build logs immediately
-         if current_log_error_state not in ['Auth Failed JWT', 'Not Found', 'Request Error', 'Poller Unexpected Error']:
-             try:
-                 run_logs_content = fetch_logs(rid, "run")
-                 if run_logs_content.startswith("ERROR_"):
-                      current_log_error_state = run_logs_content.split(':')[0].replace('ERROR_', '') # Overwrite if run log error is more specific
-                 elif current_log_error_state == 'No Logs Yet' and run_logs_content.strip() == "No run logs found yet.":
-                      pass # Keep 'No Logs Yet'
-                 elif current_log_error_state == 'No Logs Yet' and run_logs_content.strip() != "No run logs found yet.":
-                      current_log_error_state = 'None' # Logs started appearing
-
-             except Exception as e:
-                 run_logs_content = f"Poller Unexpected error fetching run logs: {e}"
-                 if current_log_error_state == 'None' or current_log_error_state == 'No Logs Yet':
-                     current_log_error_state = 'Poller Unexpected Error'
-                 print(run_logs_content)
-         else:
-              run_logs_content = f"Skipped fetching run logs due to build log error state: {current_log_error_state}"
-
-
-         try:
-            iframe_status_ok = check_iframe(url)
-            if not iframe_status_ok and current_log_error_state == 'None':
-                 # If iframe isn't OK but no log error, assume it's still building or subtle issue
-                 current_log_error_state = 'Iframe Not Ready'
-            elif iframe_status_ok and current_log_error_state == 'Iframe Not Ready':
-                 current_log_error_state = 'None' # Iframe is now ready
-
-         except Exception as e:
-            iframe_status_ok = False
-            if current_log_error_state == 'None':
-                current_log_error_state = 'Iframe Check Error'
-            print(f"Poller error checking iframe {url}: {e}")
+         current_project_state['logs']['build'] = build_logs_content
+         current_project_state['logs']['run'] = run_logs_content
+         current_project_state['iframe_ok'] = iframe_status_ok
+         current_project_state['log_error_state'] = new_log_error_state
+         # Do NOT update chat_history or status_message from the poller here.
+         # The orchestrator needs to react to the *state* changes, not manage the UI messages directly from poller.
+
     else:
-        # If not polling based on orchestration or status, return existing state/placeholders
-        print(f"Poller skipping log fetch (Status: {current_project_state.get('status')}, Task: {current_project_state.get('current_task')})")
-        pass # Keep existing logs/status
+        # If repo_id is not set in state, cannot poll logs
+         build_logs_content = current_project_state['logs'].get('build', 'Space name not set, cannot poll logs.')
+         run_logs_content = current_project_state['logs'].get('run', 'Space name not set, cannot poll logs.')
+         iframe_status_ok = False
+         current_log_error_state = current_project_state.get('log_error_state', 'Space Not Set')
+         print("Poller skipping log fetch because repo_id is not set in state.")
 
 
+    # Reconstruct preview HTML based on the updated state
     preview_html = (
-        f'<iframe src="{url}" width="100%" height="500px" allow="accelerometer; ambient-light-sensor; autoplay; camera; gyroscope; hid; fullscreen; illustration; xr-spatial-tracking; sync-xhr;" sandbox="allow-forms allow-modals allow-popups allow-popups-to-escape-sandbox allow-same-origin allow-scripts allow-storage-access-by-user-activation allow-top-navigation-by-user-activation" frameborder="0"></iframe>'
-        if iframe_status_ok else
-        f"<p style='color:red;'>⚠️ App preview not loading or check failed ({url})."
-        f"{' **Authentication Error:** Ensure HF_TOKEN secret is set on the Space running this app.' if current_log_error_state in ['Auth Failed', 'Auth Failed JWT', 'Auth Failed STREAM'] else ''}"
-        f"{' **Space or logs not found yet:** Make sure the Space name is correct and wait for the build to complete.' if current_log_error_state == 'Not Found' else ''}"
-        f"{' **Timeout:** Log fetch timed out.' if current_log_error_state == 'TIMEOUT' else ''}"
-        f"{' **Iframe not ready:** Still building or error. Check logs.' if current_log_error_state == 'Iframe Not Ready' else ''}"
-        f"{' **Polling Error:** Encountered error fetching logs.' if current_log_error_state in ['Request Error', 'Poller Unexpected Error', 'Iframe Check Error'] else ''}"
-        f"{' Check build logs for errors.' if current_log_error_state not in ['Auth Failed', 'Auth Failed JWT', 'Auth Failed STREAM', 'Not Found'] else ''}"
+        f'<iframe src="{current_project_state.get("iframe_url", "")}" width="100%" height="500px" allow="accelerometer; ambient-light-sensor; autoplay; camera; gyroscope; hid; fullscreen; illustration; xr-spatial-tracking; sync-xhr;" sandbox="allow-forms allow-modals allow-popups allow-popups-to-escape-sandbox allow-same-origin allow-scripts allow-storage-access-by-user-activation allow-top-navigation-by-user-activation" frameborder="0"></iframe>'
+        if current_project_state.get('iframe_ok') else
+        f"<p style='color:red;'>⚠️ App preview not loading or check failed ({current_project_state.get('iframe_url', 'URL not set.')})."
+        f"{' **Authentication Error:** Ensure HF_TOKEN secret is set on the Space running this app.' if current_project_state.get('log_error_state') in ['Auth Failed', 'Auth Failed JWT', 'Auth Failed STREAM'] else ''}"
+        f"{' **Space or logs not found yet:** Make sure the Space name is correct and wait for the build to complete.' if current_project_state.get('log_error_state') == 'Not Found' else ''}"
+        f"{' **Timeout:** Log fetch timed out.' if current_project_state.get('log_error_state') == 'TIMEOUT' else ''}"
+        f"{' **Iframe not ready:** Still building or error. Check logs.' if current_project_state.get('log_error_state') == 'Iframe Not Ready' else ''}"
+        f"{' **Polling Error:** Encountered error fetching logs.' if current_project_state.get('log_error_state') in ['Request Error', 'Poller Unexpected Error', 'Iframe Check Error', 'HTTP Error'] else ''}"
+        f"{' **Space Name Not Set:** Cannot fetch logs or check iframe.' if current_project_state.get('log_error_state') == 'Space Not Set' else ''}"
+        f"{' Check build logs for details.' if current_project_state.get('log_error_state') not in ['Auth Failed', 'Auth Failed JWT', 'Auth Failed STREAM', 'Not Found', 'Space Not Set'] else ''}" # Suggest checking logs if no specific critical error
         f"</p>"
     )
 
-    # Update the project state with the latest info from the poller
-    # Only update if the state exists and matches the current repo_id being polled
-    if current_project_state is not None and current_project_state.get('repo_id') == rid:
-         current_project_state['logs']['build'] = build_logs_content
-         current_project_state['logs']['run'] = run_logs_content
-         current_project_state['iframe_ok'] = iframe_status_ok
-         current_project_state['log_error_state'] = current_log_error_state
-         # Do NOT update chat_history or status_message from the poller here,
-         # as the orchestrator handles those for better flow control.
-         # The UI will display the *state's* status_message and chat_history.
 
-    return build_logs_content, run_logs_content, preview_html, current_project_state # Return updated state
+    # Return the updated state along with the UI outputs derived *from* the state
+    return (
+        build_logs_content, # Build logs output
+        run_logs_content, # Run logs output
+        preview_html,     # Preview HTML output
+        current_project_state # Updated project state
+    )
 
 
 # --- MAIN HANDLER (Called by Gradio) ---
 # Updated signature to include login state explicitly and project state first
+# Note: Gradio 4.x might prefer profile and token as the first two args when auto-injecting.
+# Let's stick to the recommended Gradio pattern where auto-injected params come first.
 def handle_user_message(
-    profile: gr.OAuthProfile | None, # (1) - from login_btn (first in inputs)
-    oauth_token: gr.OAuthToken | None, # (2) - from login_btn (first in inputs)
-    project_state: dict | None, # (3) - from gr.State
+    profile: gr.OAuthProfile | None, # (1) - Auto-injected by Gradio
+    oauth_token: gr.OAuthToken | None, # (2) - Auto-injected by Gradio
+    project_state: dict | None, # (3) - from gr.State component
     user_input: str,     # (4)
     sdk_choice: str,     # (5)
     gemini_api_key: str, # (6)
@@ -1077,15 +1072,17 @@ def handle_user_message(
     temperature: float,  # (9)
     max_output_tokens: int, # (10)
 ):
-    # Now the function expects 10 parameters.
-    # The inputs list will be structured to match this.
+    # Now the function signature has 10 parameters.
+    # The inputs list will exclude login_btn to allow auto-injection.
 
     # Initialize or update project_state based on inputs and previous state
     # Decide whether to start a new project or continue
+    # A new project starts if state is None, or if requirements change
     is_new_project_request = (
         project_state is None or # First run
-        project_state.get('status') in ['Complete', 'Failed'] or # Previous project finished
         (user_input is not None and user_input.strip() != "" and user_input.strip() != project_state.get('requirements', '').strip()) # New requirements provided
+        # We could also check if status is 'Complete' or 'Failed' to auto-reset, but letting
+        # the user input change trigger the reset is more explicit.
     )
 
     if is_new_project_request:
@@ -1098,7 +1095,7 @@ def handle_user_message(
         sdk_version = get_sdk_version(sdk_choice)
 
         project_state = {
-            'requirements': user_input,
+            'requirements': user_input.strip() if user_input else '', # Store cleaned requirements
             'plan': '',
             'files': {},
             'logs': {'build': '', 'run': ''},
@@ -1116,60 +1113,79 @@ def handle_user_message(
             'iframe_ok': False,
             'log_error_state': 'None',
         }
-        project_state['chat_history'].append({"role": "user", "content": user_input}) # Add initial user requirement
+        if user_input and user_input.strip():
+             project_state['chat_history'].append({"role": "user", "content": user_input.strip()}) # Add initial user requirement
 
     else:
-        # Continue existing project state - just update potentially changed inputs
+        # Continue existing project state
         print("Continuing existing project...")
-        # history is part of project_state
-        # project_state['requirements'] remains the requirement that started the project
-        project_state['sdk_choice'] = sdk_choice # Allow changing SDK mid-project? Maybe not intended.
-        project_state['main_app_file'] = "app.py" if sdk_choice == "gradio" else "streamlit_app.py" # Update main file name
-        project_state['sdk_version'] = get_sdk_version(sdk_choice) # Update version
+        # history is part of project_state - no need to copy
+        # project_state['requirements'] stays the same requirement that started the project
+        # Update other settings that could change mid-project (though maybe shouldn't)
+        project_state['sdk_choice'] = sdk_choice
+        project_state['main_app_file'] = "app.py" if sdk_choice == "gradio" else "streamlit_app.py"
+        project_state['sdk_version'] = get_sdk_version(sdk_choice)
         clean_space_name = space_name.strip() if space_name else ''
-        project_state['repo_id'] = f"{profile.username}/{clean_space_name}" if profile and clean_space_name else ''
-        project_state['iframe_url'] = f"https://huggingface.co/spaces/{profile.username}/{clean_space_name}" if profile and clean_space_name else ''
-        # Do NOT change task, status, attempts, logs, feedback, plan, files here.
-        # These are managed by the orchestration loop.
-        # The user_input might trigger a new plan if the Orchestrator is made to handle that.
-        # For now, a new user_input starts a new project (handled by is_new_project_request).
+        # Only update repo_id/iframe_url if the name changes (and state is not 'Not Started')
+        if project_state.get('repo_id') != f"{profile.username}/{clean_space_name}" and project_state.get('status') != 'Not Started' :
+             # This case should ideally be handled by the poller's state reset, but double check.
+             print(f"Handle_user_message detected space name change mid-project from {project_state.get('repo_id')} to {profile.username}/{clean_space_name}. Forcing project reset.")
+             # Re-initialize as a new project
+             return handle_user_message(profile, oauth_token, None, user_input, sdk_choice, gemini_api_key, space_name, grounding_enabled, temperature, max_output_tokens)
+        elif profile and clean_space_name:
+            project_state['repo_id'] = f"{profile.username}/{clean_space_name}"
+            project_state['iframe_url'] = f"https://huggingface.co/spaces/{profile.username}/{clean_space_name}"
+
+        # If status is 'Not Started' and user sends input again (maybe after fixing something),
+        # set task to START to begin orchestration.
+        if project_state.get('status') == 'Not Started' and user_input and user_input.strip() and project_state.get('current_task') is None:
+             print("Status was 'Not Started', user sent input. Starting orchestration.")
+             project_state['current_task'] = 'START'
+             project_state['status'] = 'In Progress'
+             project_state['status_message'] = 'Initializing...'
 
 
     # Validation Checks (using received profile/token)
-    if not profile or not oauth_token or not hasattr(oauth_token, 'token') or not oauth_token.token:
-        error_msg = "⚠️ Please log in first via the Hugging Face button."
-        if not project_state['chat_history'] or project_state['chat_history'][-1].get("content") != error_msg:
-            project_state['chat_history'].append({"role":"assistant","content":error_msg})
-        project_state['status'] = 'Failed' # Mark as failed due to login
-        project_state['status_message'] = "Login required."
-        project_state['current_task'] = 'FINISHED' # End orchestration
+    validation_error = None
+    validation_message = None
 
+    if not profile or not oauth_token or not hasattr(oauth_token, 'token') or not oauth_token.token:
+        validation_error = "Login required."
+        validation_message = "⚠️ Please log in first via the Hugging Face button."
     elif not space_name or not space_name.strip():
-        msg = "⚠️ Please enter a Space name."
-        if not project_state['chat_history'] or project_state['chat_history'][-1].get("content") != msg:
-            project_state['chat_history'].append({"role":"assistant","content":msg})
-        project_state['status'] = 'Failed' # Mark as failed due to missing space name
-        project_state['status_message'] = "Space name required."
-        project_state['current_task'] = 'FINISHED' # End orchestration
-
+        validation_error = "Space name required."
+        validation_message = "⚠️ Please enter a Space name."
     elif not gemini_api_key:
-         error_msg = "⚠️ Please provide your Gemini API Key."
-         if not project_state['chat_history'] or project_state['chat_history'][-1].get("content") != error_msg:
-            project_state['chat_history'].append({"role":"assistant","content":error_msg})
-         project_state['status'] = 'Failed' # Mark as failed due to missing API key
-         project_state['status_message'] = "API Key required."
-         project_state['current_task'] = 'FINISHED' # End orchestration
-
+         validation_error = "API Key required."
+         validation_message = "⚠️ Please provide your Gemini API Key."
     elif not user_input or user_input.strip() == "":
-        error_msg = "Please enter requirements for the application."
-        if not project_state['chat_history'] or project_state['chat_history'][-1].get("content") != error_msg:
-            project_state['chat_history'].append({"role":"assistant","content":error_msg})
-        # Don't fail immediately, wait for requirements
-        project_state['status_message'] = "Waiting for prompt."
-        project_state['status'] = 'Not Started' # Explicitly set state while waiting
-
-
-    # If validation passed and status is still 'In Progress', run orchestration
+        validation_error = "Requirements required."
+        validation_message = "Please enter requirements for the application."
+
+
+    # If there's a validation error, update state and return early
+    if validation_error:
+         if project_state.get('status') != 'Failed' or project_state.get('status_message') != validation_message:
+             # Avoid adding duplicate error messages
+             project_state['status'] = 'Failed' # Mark as failed due to validation
+             project_state['status_message'] = validation_message
+             project_state['current_task'] = 'FINISHED' # End orchestration
+             if not project_state['chat_history'] or project_state['chat_history'][-1].get("content") != validation_message:
+                project_state['chat_history'].append({"role":"assistant","content":validation_message})
+
+         # Return current state and UI outputs
+         return (
+             project_state,
+             project_state.get('chat_history', []),
+             project_state['logs'].get('build', ''),
+             project_state['logs'].get('run', ''),
+             f"<p style='color:red;'>{validation_message}</p>", # Simple error preview
+             project_state.get('status_message', 'Validation Failed')
+         )
+
+
+    # If validation passed and status is 'In Progress' or being set to 'In Progress' ('START' task)
+    # Note: orchestrate_development handles the 'START' -> 'In Progress' transition internally
     if project_state.get('status') == 'In Progress':
         client = genai.Client(api_key=gemini_api_key)
 
@@ -1181,10 +1197,11 @@ def handle_user_message(
         )
 
         # Run *one step* of orchestration. orchestrate_development updates project_state directly.
-        # The function is designed to be called repeatedly.
+        # The function is designed to be called repeatedly by subsequent handle_user_message calls
+        # triggered by the poller implicitly updating state and thus triggering retries.
         try:
             orchestrate_development(
-                 client, project_state, cfg, oauth_token.token if oauth_token else None
+                 client, project_state, cfg, oauth_token.token # Pass the token for pushing
             )
         except Exception as e:
              # Catch unexpected errors during orchestration step
@@ -1213,8 +1230,9 @@ def handle_user_message(
          f"{' **Space or logs not found yet:** Make sure the Space name is correct and wait for the build to complete.' if project_state.get('log_error_state') == 'Not Found' else ''}"
          f"{' **Timeout:** Log fetch timed out.' if project_state.get('log_error_state') == 'TIMEOUT' else ''}"
          f"{' **Iframe not ready:** Still building or error. Check logs.' if project_state.get('log_error_state') == 'Iframe Not Ready' else ''}"
-         f"{' **Polling Error:** Encountered error fetching logs.' if project_state.get('log_error_state') in ['Request Error', 'Poller Unexpected Error', 'Iframe Check Error'] else ''}"
-         f"{' Check build logs for errors.' if project_state.get('log_error_state') in ['None', 'No Logs Yet', 'Iframe Not Ready'] else ''}" # Suggest checking logs if no specific error or just waiting
+         f"{' **Polling Error:** Encountered error fetching logs.' if project_state.get('log_error_state') in ['Request Error', 'Poller Unexpected Error', 'Iframe Check Error', 'HTTP Error'] else ''}"
+         f"{' **Space Name Not Set:** Cannot fetch logs or check iframe.' if project_state.get('log_error_state') == 'Space Not Set' else ''}"
+         f"{' Check build logs for details.' if project_state.get('log_error_state') not in ['Auth Failed', 'Auth Failed JWT', 'Auth Failed STREAM', 'Not Found', 'Space Not Set'] else ''}" # Suggest checking logs if no specific critical error
          f"</p>"),
         project_state.get('status_message', '...'), # Status message output (get from state)
     )
@@ -1238,17 +1256,19 @@ with gr.Blocks(title="HF Space Auto‑Builder (Team AI)") as demo:
             models_md = gr.Markdown()
 
             # Use api_name to make these callable via API if needed, but not strictly necessary for UI load
+            # Ensure these load on initial page load
             demo.load(show_profile, inputs=None, outputs=status_md) # api_name="load_profile" removed
             demo.load(list_private_models, inputs=None, outputs=models_md) # api_name="load_models" removed
 
+            # Ensure these update on login button click
             login_btn.click(
                 fn=show_profile,
-                inputs=None,
+                inputs=None, # OAuth auto-injected
                 outputs=status_md
             )
             login_btn.click(
                 fn=list_private_models,
-                inputs=None,
+                inputs=None, # OAuth auto-injected
                 outputs=models_md
             )
             # --- END LOGIN FIX ---
@@ -1283,17 +1303,15 @@ with gr.Blocks(title="HF Space Auto‑Builder (Team AI)") as demo:
         preview   = gr.HTML("<p>App preview will load here when available.</p>")
 
     # --- Hidden poller (using Timer) ---
-    # Polls less frequently initially, increases if project is in progress? Or keep constant?
-    # Constant 2s is fine for demonstration.
+    # Polls every 2 seconds to update logs and preview
     log_poller = gr.Timer(value=2, active=True, render=False)
     # --- End hidden poller ---
 
     # handle_user_message is defined ABOVE this block
     # The main button and submit handler now receive and return the project_state
-    # INPUTS: gr.LoginButton outputs (profile, token) first, then state, then other UI inputs
+    # INPUTS: relies on auto-injection of profile/token, then state, then other UI inputs
     inputs_list = [
-         login_btn, # <-- Login button first (provides profile, token)
-         project_state, # <-- Pass the state in
+         project_state, # <-- Pass the state in (after auto-injected profile/token)
          user_in,       # Other UI inputs in order
          sdk_choice,
          api_key,
@@ -1301,8 +1319,9 @@ with gr.Blocks(title="HF Space Auto‑Builder (Team AI)") as demo:
          grounding,
          temp,
          max_tokens,
+         # Note: login_btn is NOT in inputs, profile/token are auto-injected
          # Note: chatbot is NOT in inputs, its history is managed via project_state['chat_history']
-    ] # Total 9 components -> should provide 2 + 1 + 7 = 10 arguments
+    ] # Total 8 components listed here, plus 2 auto-injected = 10 arguments
 
     # OUTPUTS: Must match the return values of handle_user_message
     outputs_list = [
@@ -1318,23 +1337,27 @@ with gr.Blocks(title="HF Space Auto‑Builder (Team AI)") as demo:
         fn=handle_user_message,
         inputs=inputs_list,
         outputs=outputs_list
+        # No api_name needed unless you plan to call this via API directly
     )
 
     user_in.submit(
         fn=handle_user_message,
         inputs=inputs_list,
         outputs=outputs_list
+        # No api_name needed
     )
 
     # --- Wire the poller to the update function ---
     # _update_logs_and_preview is defined ABOVE this block
-    # The poller needs login state, space name, and the current project state
+    # The poller needs login state (as tuple), space name, and the current project state
     # INPUTS: login_btn state (as tuple), space name, project state
+    # When login_btn is the *first* component in the inputs list and the function's
+    # first parameter is a tuple of the correct types, Gradio passes the outputs as a tuple.
     poller_inputs = [
-         login_btn,    # Provides (profile, token) as a tuple when put first for a single param
+         login_btn,    # Provides (profile, token) as a tuple because it's first
          space_name,
          project_state # Passes the current state value
-    ] # Total 3 components -> should provide 1 (tuple) + 1 + 1 = 3 arguments
+    ] # Total 3 components -> function receives 3 arguments (1 tuple + 1 + 1)
     # Note: _update_logs_and_preview signature is (profile_token_state, space_name, current_project_state)
 
     # OUTPUTS: Must match the return values of _update_logs_and_preview
@@ -1349,6 +1372,7 @@ with gr.Blocks(title="HF Space Auto‑Builder (Team AI)") as demo:
         fn=_update_logs_and_preview,
         inputs=poller_inputs,
         outputs=poller_outputs
+        # No api_name needed
     )
     # --- End wire poller ---
 
@@ -1360,4 +1384,5 @@ with gr.Blocks(title="HF Space Auto‑Builder (Team AI)") as demo:
 # If running on a Space, server_name='0.0.0.0' and server_port=7860 are standard.
 # If running locally, just demo.launch() or specify a port.
 if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+    # Use `show_error=True` to see the actual Python errors in the UI during development
+    demo.launch(server_name="0.0.0.0", server_port=7860, show_error=True)