Update backend/orchestrator.py

backend/orchestrator.py

@@ -8,7 +8,7 @@ import ast
 from typing import Optional, Dict, Any
 
 from .database import update_project_status, get_project
-from .tools import get_project_dir, create_file, read_file, zip_project, scrape_web, document_requirements, create_design_spec
+from .tools import get_project_dir, create_file, read_file, zip_project, scrape_web, document_requirements, create_design_spec, run_test_in_sandbox, run_bandit
 from .agents import ROLE_PROMPTS
 from models.loader import generate_with_model
 
@@ -65,16 +65,14 @@ def run_agent_chain(project_id, user_id, initial_prompt):
             "max_turns": 60,
             "current_task_index": 0,
             "error_state": None,
-            "phase": "requirements_gathering" # New state variable
+            "phase": "requirements_gathering"
         }
         
         next_agent = "business_analyst"
 
-        # --- PHASED EXECUTION LOOP ---
         while project_context["turn"] < project_context["max_turns"]:
             log_step("ORCHESTRATOR", f"Current Phase: {project_context['phase'].replace('_', ' ').title()}")
             
-            # --- PHASE 1: REQUIREMENTS GATHERING ---
             if project_context["phase"] == "requirements_gathering":
                 prompt = f"""You are the Business Analyst. The user's request is: {project_context['initial_prompt']}. 
                 Based on this, generate a list of functional requirements."""
@@ -89,7 +87,6 @@ def run_agent_chain(project_id, user_id, initial_prompt):
                     next_agent = "ux_ui_designer"
                 else: raise ValueError("Business Analyst failed to provide a valid output.")
             
-            # --- PHASE 2: UX/UI DESIGN ---
             elif project_context["phase"] == "design_phase":
                 prompt = f"""You are the UX/UI Designer. The documented requirements are: {project_context['requirements']}.
                 Create a detailed user interface design plan based on these requirements."""
@@ -101,7 +98,6 @@ def run_agent_chain(project_id, user_id, initial_prompt):
                     log_step("UX/UI DESIGNER", "Design specification created.", str(designer_output["design"]))
                     create_design_spec(project_dir, json.dumps(designer_output, indent=2))
                     
-                    # Now, a Planner is needed to break down the design into coding tasks
                     log_step("ORCHESTRATOR", "Design is complete. Briefing the Planner to create a coding plan.")
                     planner_prompt = f"Project brief: {project_context['initial_prompt']}\nRequirements: {project_context['requirements']}\nDesign: {project_context['design']}\nCreate a step-by-step plan for the developers. Assign tasks to 'worker_backend_coder', 'worker_front_end_coder', or 'worker_tester'."
                     planner_response = generate_with_model("planner", planner_prompt)
@@ -111,10 +107,9 @@ def run_agent_chain(project_id, user_id, initial_prompt):
                     project_context["detailed_plan"] = planner_output["plan"]
                     log_step("PLANNER", "Coding plan created.", f"Plan: {project_context['detailed_plan']}")
                     project_context["phase"] = "development"
-                    next_agent = "manager" # Return control to manager for execution
+                    next_agent = "manager"
                 else: raise ValueError("UX/UI Designer failed to provide a valid output.")
 
-            # --- PHASE 3: DEVELOPMENT AND DEBUGGING ---
             elif project_context["phase"] == "development":
                 if project_context["current_task_index"] >= len(project_context["detailed_plan"]):
                     log_step("ORCHESTRATOR", "All planned tasks completed. Instructing Manager to finish project.")
@@ -148,12 +143,10 @@ def run_agent_chain(project_id, user_id, initial_prompt):
                 log_step(next_agent.replace("_", " ").title(), f"Decided to perform action: `{action}` with args: `{args}`")
 
                 if action == "delegate_coder":
-                    # This needs to be a bit smarter to delegate to the right coder
                     next_agent = current_plan_step["responsible_agent"]
                     project_context["current_task"] = args[0]
                     project_context["last_action_result"] = f"Task delegated to {next_agent.replace('_', ' ')}."
-
-                # --- Handle new worker agent actions ---
+                
                 elif action in ["write_code", "write_test", "task_complete"]:
                     if action == "write_code" and len(args) >= 2: create_file(project_dir, args[0], args[1])
                     if action == "write_test" and len(args) >= 2: create_file(project_dir, args[0], args[1])
@@ -161,6 +154,37 @@ def run_agent_chain(project_id, user_id, initial_prompt):
                     project_context["last_action_result"] = f"Worker finished task and saved to file. Moving to next step in the plan."
                     project_context["current_task_index"] += 1
                     next_agent = "manager"
+                
+                elif action == "run_test":
+                    if not args: raise ValueError("Missing file path for run_test.")
+                    test_file_path = args[0]
+                    test_output = run_test_in_sandbox(project_dir, test_file_path)
+                    
+                    project_context["last_action_result"] = f"Test run for {test_file_path}:\n{test_output}"
+                    log_step("TESTER", "Test run complete.", test_output)
+                    
+                    if "Sandbox Execution Error" in test_output or "Traceback (most recent call last)" in test_output:
+                        log_step("ORCHESTRATOR", "Tests failed. Handing over to Debugger.")
+                        next_agent = "debugger"
+                        project_context["error_state"] = test_output
+                    else:
+                        log_step("ORCHESTRATOR", "Tests passed. Continuing with the plan.")
+                        project_context["current_task_index"] += 1
+                        next_agent = "manager"
+                
+                elif action == "run_bandit":
+                    bandit_report = run_bandit(project_dir)
+                    project_context["last_action_result"] = f"Bandit Analysis Report:\n{bandit_report}"
+                    log_step("CODE ANALYST", "Static analysis complete.", bandit_report)
+                    
+                    if "issues" in bandit_report and json.loads(bandit_report)["results"]:
+                        log_step("ORCHESTRATOR", "Bandit found security issues. Handing over to Debugger.")
+                        next_agent = "debugger"
+                        project_context["error_state"] = bandit_report
+                    else:
+                        log_step("ORCHESTRATOR", "No issues found by Bandit. Continuing with the plan.")
+                        project_context["current_task_index"] += 1
+                        next_agent = "manager"
 
                 elif action == "propose_fix":
                     fix_description = args[2]
@@ -182,7 +206,7 @@ def run_agent_chain(project_id, user_id, initial_prompt):
                         url = args[0]
                         scraped_content = scrape_web(url)
                         project_context["last_action_result"] = f"Scraping result from {url}:\n\n{scraped_content}"
-                    continue # Do not increment turn
+                    continue
                 
                 elif action == "finish_project":
                     log_step("MANAGER", "Project is ready for final packaging.")