feat: Implement real-time progress updates for test generation

This commit introduces a real-time progress update feature to the UI. When a user clicks "Generate & Execute Test", the backend now streams progress updates to the frontend, which are displayed in the "Execution Results" panel.

Key changes:
- Modified the agentic workflow in `robot_generator.py` to be a generator that yields status updates at each step of the process.
- Refactored the `/generate-and-run` endpoint in `main.py` to use Server-Sent Events (SSE) with a heartbeat mechanism. The heartbeat prevents proxy buffering issues and ensures real-time delivery of messages.
- The synchronous agentic workflow is run in a separate thread to avoid blocking the async server. A queue is used to communicate between the thread and the main event loop.
- Updated the frontend JavaScript in `index.html` to handle the SSE stream and display the progress messages in a running log format in the UI.

Author: google-labs-jules[bot]

backend/main.py

@@ -9,6 +9,9 @@
 from dotenv import load_dotenv
 from fastapi.middleware.cors import CORSMiddleware
 import json
+import asyncio
+from queue import Queue, Empty
+from threading import Thread
 
 # Import the new agentic workflow orchestrator
 from backend.robot_generator import run_agentic_workflow
@@ -37,29 +40,72 @@ class Query(BaseModel):
 )
 
 # --- Main Endpoint ---
-async def stream_generate_and_run(user_query: str, model_name: str):
-    """Generator function to stream logs and results."""
+def run_workflow_in_thread(queue, user_query, model_provider, model_name):
+    """Runs the synchronous agentic workflow and puts results in a queue."""
+    try:
+        for event in run_agentic_workflow(user_query, model_provider, model_name):
+            queue.put(event)
+    except Exception as e:
+        logging.error(f"Exception in workflow thread: {e}")
+        queue.put({"status": "error", "message": f"Workflow thread failed: {e}"})
 
-    # Stage 1: Generating Code
-    yield f"data: {json.dumps({'stage': 'generation', 'status': 'running', 'message': 'AI agents are analyzing your query...'})}\n\n"
+async def stream_generate_and_run(user_query: str, model_name: str):
+    """
+    Generator function that streams logs and results, with a heartbeat
+    to prevent timeouts and buffer flushing issues.
+    """
+    robot_code = None
+    model_provider = os.getenv("MODEL_PROVIDER", "online").lower()
+    q = Queue()
+
+    # Run the synchronous workflow in a separate thread
+    workflow_thread = Thread(
+        target=run_workflow_in_thread,
+        args=(q, user_query, model_provider, model_name)
+    )
+    workflow_thread.start()
+
+    # --- Stage 1: Generating Code ---
+    while workflow_thread.is_alive():
+        try:
+            event = q.get_nowait()
+            event_data = {'stage': 'generation', **event}
+            yield f"data: {json.dumps(event_data)}\n\n"
+
+            if event.get("status") == "complete" and "robot_code" in event:
+                robot_code = event["robot_code"]
+                # Generation is done, we can break this loop and move to execution
+                workflow_thread.join() # Ensure thread is cleaned up
+                break
+            elif event.get("status") == "error":
+                logging.error(f"Error during code generation: {event.get('message')}")
+                workflow_thread.join()
+                return
+        except Empty:
+            # Send a heartbeat comment to keep the connection open
+            yield ": heartbeat\n\n"
+            await asyncio.sleep(1)
+
+    # In case the thread finished but we didn't get the code
+    if not robot_code:
+        # Check the queue one last time
+        while not q.empty():
+            event = q.get_nowait()
+            event_data = {'stage': 'generation', **event}
+            yield f"data: {json.dumps(event_data)}\n\n"
+            if event.get("status") == "complete" and "robot_code" in event:
+                robot_code = event["robot_code"]
+            elif event.get("status") == "error":
+                return # Error was already sent
 
-    try:
-        model_provider = os.getenv("MODEL_PROVIDER", "online").lower()
-        robot_code = run_agentic_workflow(
-            natural_language_query=user_query,
-            model_provider=model_provider,
-            model_name=model_name
-        )
         if not robot_code:
-            raise Exception("Agentic workflow failed to generate Robot Framework code.")
+            final_error_message = "Agentic workflow finished without generating code."
+            logging.error(final_error_message)
+            yield f"data: {json.dumps({'stage': 'generation', 'status': 'error', 'message': final_error_message})}\n\n"
+            return
 
-        yield f"data: {json.dumps({'stage': 'generation', 'status': 'complete', 'message': 'Code generation complete.', 'robot_code': robot_code})}\n\n"
-    except Exception as e:
-        logging.error(f"Error during code generation: {e}")
-        yield f"data: {json.dumps({'stage': 'generation', 'status': 'error', 'message': str(e)})}\n\n"
-        return
 
-    # Stage 2: Docker Execution
+    # --- Stage 2: Docker Execution ---
     run_id = str(uuid.uuid4())
     robot_tests_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), '..', 'robot_tests', run_id)
     os.makedirs(robot_tests_dir, exist_ok=True)

backend/robot_generator.py

@@ -288,53 +288,65 @@ def agent_code_validator(code: str, model_provider: str, model_name: str) -> Val
 
 # --- Orchestrator ---
 
-def run_agentic_workflow(natural_language_query: str, model_provider: str, model_name: str) -> Optional[str]:
+def run_agentic_workflow(natural_language_query: str, model_provider: str, model_name: str):
     """
     Orchestrates the multi-agent workflow to generate Robot Framework code,
-    including a self-correction loop.
+    yielding progress updates and the final code.
     """
     logging.info("--- Starting Multi-Agent Workflow ---")
+    yield {"status": "running", "message": "Starting agentic workflow..."}
     MAX_ATTEMPTS = 3
 
     # Configure online provider if used
     if model_provider == "online":
         api_key = os.getenv("GEMINI_API_KEY")
         if not api_key:
             logging.error("Orchestrator: GEMINI_API_KEY not found for online provider.")
-            raise ValueError("GEMINI_API_KEY not found in environment variables.")
+            yield {"status": "error", "message": "GEMINI_API_KEY not found."}
+            return
         genai.configure(api_key=api_key)
 
-    # The initial query for the first attempt
     current_query = natural_language_query
 
     for attempt in range(MAX_ATTEMPTS):
         logging.info(f"--- Attempt {attempt + 1} of {MAX_ATTEMPTS} ---")
+        yield {"status": "running", "message": f"Starting attempt {attempt + 1}/{MAX_ATTEMPTS}..."}
 
         # Agent 1: Plan
+        yield {"status": "running", "message": "Agent 1/4: Planning test steps..."}
         planned_steps = agent_step_planner(current_query, model_provider, model_name)
         if not planned_steps:
             logging.error("Orchestrator: Step Planner failed. Aborting.")
-            return None
+            yield {"status": "error", "message": "Failed to generate a test plan."}
+            return
+        yield {"status": "running", "message": "Agent 1/4: Test step planning complete."}
 
         # Agent 2: Identify Locators
+        yield {"status": "running", "message": "Agent 2/4: Identifying UI element locators..."}
         located_steps = agent_element_identifier(planned_steps, model_provider, model_name)
         if not located_steps:
             logging.error("Orchestrator: Element Identifier failed. Aborting.")
-            return None
+            yield {"status": "error", "message": "Failed to identify UI element locators."}
+            return
+        yield {"status": "running", "message": "Agent 2/4: UI element locator identification complete."}
 
         # Agent 3: Assemble Code
+        yield {"status": "running", "message": "Agent 3/4: Assembling Robot Framework code..."}
         robot_code = agent_code_assembler(located_steps, natural_language_query)
+        yield {"status": "running", "message": "Agent 3/4: Code assembly complete."}
 
         # Agent 4: Validate
+        yield {"status": "running", "message": "Agent 4/4: Validating generated code..."}
         validation = agent_code_validator(robot_code, model_provider, model_name)
+        yield {"status": "running", "message": "Agent 4/4: Code validation complete."}
 
         if validation.valid:
             logging.info("Code validation successful. Workflow complete.")
-            logging.info("--- Multi-Agent Workflow Complete ---")
-            return robot_code
+            yield {"status": "complete", "robot_code": robot_code, "message": "Code generation successful."}
+            return
         else:
             logging.warning(f"Code validation failed. Reason: {validation.reason}")
-            # Prepare for the next attempt by creating a corrective query
+            yield {"status": "running", "message": f"Validation failed: {validation.reason}. Attempting self-correction..."}
             current_query = f"""
             The previous attempt to generate a test plan failed validation.
             The user's original query was: "{natural_language_query}"
@@ -349,5 +361,4 @@ def run_agentic_workflow(natural_language_query: str, model_provider: str, model
             logging.info("Attempting self-correction...")
 
     logging.error("Orchestrator: Failed to generate valid code after multiple attempts.")
-    logging.info("--- Multi-Agent Workflow Failed ---")
-    return None
+    yield {"status": "error", "message": "Failed to generate valid code after multiple attempts."}

frontend/index.html

@@ -626,8 +626,8 @@ <h1>AI Test Automation Platform</h1>
                 }
 
                 setButtonLoading(true);
-                robotCodeEl.innerHTML = '';
-                executionLogsEl.innerHTML = '';
+                robotCodeEl.innerHTML = '<div class="empty-state"><div class="empty-state-icon">⚡</div><p>Generated Robot Framework code will appear here</p></div>';
+                executionLogsEl.innerHTML = ''; // Clear previous logs
                 downloadBtn.style.display = 'none';
                 robotCodeContent = '';
 
@@ -687,19 +687,40 @@ <h1>AI Test Automation Platform</h1>
                 updateStatus('processing', data.message);
 
                 if (data.stage === 'generation') {
-                    if (data.status === 'complete' && data.robot_code) {
+                    // Display running logs for generation stage
+                    if (data.status === 'running') {
+                        const logEntry = document.createElement('div');
+                        logEntry.textContent = `[${new Date().toLocaleTimeString()}] ${data.message}`;
+                        executionLogsEl.appendChild(logEntry);
+                        executionLogsEl.scrollTop = executionLogsEl.scrollHeight; // Auto-scroll
+                    } else if (data.status === 'complete' && data.robot_code) {
                         robotCodeContent = data.robot_code;
                         robotCodeEl.textContent = robotCodeContent;
                         downloadBtn.style.display = 'inline-flex';
+                        // Clear generation logs and prepare for execution logs
+                        executionLogsEl.innerHTML = '<div class="empty-state"><div class="empty-state-icon">📋</div><p>Test execution logs will appear here</p></div>';
                     } else if (data.status === 'error') {
                         updateStatus('error', 'Generation failed');
                         robotCodeEl.innerHTML = `<div class="empty-state"><div class="empty-state-icon">⚠️</div><p>${data.message}</p></div>`;
+                        const errorEntry = document.createElement('div');
+                        errorEntry.style.color = 'var(--error)';
+                        errorEntry.textContent = `[${new Date().toLocaleTimeString()}] ERROR: ${data.message}`;
+                        executionLogsEl.appendChild(errorEntry);
                         hideStatus();
                     }
                 } else if (data.stage === 'execution') {
-                    if (data.status === 'complete' && data.result) {
+                    // When execution starts, clear the placeholder and show the first real log
+                    if (executionLogsEl.querySelector('.empty-state')) {
+                        executionLogsEl.innerHTML = '';
+                    }
+                    if (data.status === 'running') {
+                        const logEntry = document.createElement('div');
+                        logEntry.textContent = `[${new Date().toLocaleTimeString()}] ${data.message}`;
+                        executionLogsEl.appendChild(logEntry);
+                        executionLogsEl.scrollTop = executionLogsEl.scrollHeight;
+                    } else if (data.status === 'complete' && data.result) {
                         const logs = data.result.logs || 'No execution logs available';
-                        executionLogsEl.textContent = logs;
+                        executionLogsEl.textContent = logs; // Replace with final, full logs
                         if (logs.includes('PASSED')) {
                             updateStatus('success', 'Test passed');
                         } else {
@@ -708,7 +729,10 @@ <h1>AI Test Automation Platform</h1>
                         hideStatus();
                     } else if (data.status === 'error') {
                         updateStatus('error', 'Execution failed');
-                        executionLogsEl.textContent = data.message;
+                        const errorEntry = document.createElement('div');
+                        errorEntry.style.color = 'var(--error)';
+                        errorEntry.textContent = `[${new Date().toLocaleTimeString()}] ERROR: ${data.message}`;
+                        executionLogsEl.appendChild(errorEntry);
                         hideStatus();
                     }
                 }