Update graph.py

graph.py

@@ -1,4 +1,4 @@
-# graph.py - Enhanced with better loop control and cost tracking
+# graph.py - Fixed version with proper state handling for concurrent updates
 
 import json
 import re
@@ -8,7 +8,7 @@ import uuid
 import shutil
 import zipfile
 import operator
-from typing import TypedDict, List, Dict, Optional, Annotated
+from typing import TypedDict, List, Dict, Optional, Annotated, Any
 from datetime import datetime
 from langchain_openai import ChatOpenAI
 from langgraph.graph import StateGraph, END
@@ -76,11 +76,36 @@ class AgentState(TypedDict):
     execution_path: Annotated[List[str], operator.add]
     rework_cycles: int
     max_loops: int
-    status_update: str
-    # NEW: For real-time cost tracking
+    # Use Annotated with operator.add for fields that multiple agents might update
+    status_updates: Annotated[List[Dict[str, str]], operator.add]  # Changed from status_update
     current_cost: float
     budget_exceeded: bool
-
+    # Add other fields that might have concurrent updates
+    pragmatistReport: Optional[Dict]
+    governanceReport: Optional[Dict]
+    complianceReport: Optional[Dict]
+    observerReport: Optional[Dict]
+    knowledgeInsights: Optional[Dict]
+
+# Helper to get latest status
+def get_latest_status(state: AgentState) -> str:
+    """Get the most recent status update from the list"""
+    updates = state.get('status_updates', [])
+    if updates and isinstance(updates, list):
+        # Get the last update's status value
+        for update in reversed(updates):
+            if isinstance(update, dict) and 'status' in update:
+                return update['status']
+            elif isinstance(update, str):
+                return update
+    return "Processing..."
+
+# Helper to add status update
+def add_status_update(node_name: str, status: str) -> Dict[str, Any]:
+    """Create a status update entry"""
+    return {
+        "status_updates": [{"node": node_name, "status": status, "timestamp": datetime.utcnow().isoformat()}]
+    }
 
 # --- LLM ---
 llm = ChatOpenAI(model="gpt-4o", temperature=0.5, max_retries=3, request_timeout=60)
@@ -187,7 +212,7 @@ def parse_json_from_llm(llm_output: str) -> Optional[dict]:
         except Exception as e:
             logger.debug(f"json.loads still failed after cleanup: {e}")
 
-    # nothing parsed — log preview and return None
+    # nothing parsed – log preview and return None
     logger.error("parse_json_from_llm failed to parse LLM output. LLM output preview (200 chars): %s", text[:200].replace("\n","\\n"))
     return None
 
@@ -368,8 +393,15 @@ def run_triage_agent(state: AgentState):
     response = llm.invoke(prompt)
     content = getattr(response, "content", "") or ""
     if 'greeting' in content.lower():
-        return {"draftResponse": "Hello! How can I help?", "execution_path": ["Triage"], "status_update": "Greeting"}
-    return {"execution_path": ["Triage"], "status_update": "Task detected"}
+        return {
+            "draftResponse": "Hello! How can I help?", 
+            "execution_path": ["Triage"],
+            **add_status_update("Triage", "Greeting")
+        }
+    return {
+        "execution_path": ["Triage"],
+        **add_status_update("Triage", "Task detected")
+    }
 
 def run_planner_agent(state: AgentState):
     log.info("--- PLANNER ---")
@@ -378,7 +410,11 @@ def run_planner_agent(state: AgentState):
     response = llm.invoke(prompt)
     plan_data = parse_json_from_llm(getattr(response, "content", "") or "")
     if not plan_data:
-        return {"pmPlan": {"error": "Planning failed"}, "execution_path": path, "status_update": "Error"}
+        return {
+            "pmPlan": {"error": "Planning failed"}, 
+            "execution_path": path,
+            **add_status_update("Planner", "Error")
+        }
     
     calls = plan_data.get('estimated_llm_calls_per_loop', 3)
     cost_per_loop = (calls * AVG_TOKENS_PER_CALL) * ((GPT4O_INPUT_COST_PER_1K_TOKENS + GPT4O_OUTPUT_COST_PER_1K_TOKENS) / 2)
@@ -392,14 +428,22 @@ def run_planner_agent(state: AgentState):
         plan_data.setdefault('experiment_type', detection.get('artifact_type'))
         plan_data.setdefault('experiment_goal', state.get('userInput',''))
     
-    return {"pmPlan": plan_data, "execution_path": path, "status_update": "Plan created"}
+    return {
+        "pmPlan": plan_data, 
+        "execution_path": path,
+        **add_status_update("Planner", "Plan created")
+    }
 
 def run_memory_retrieval(state: AgentState):
     log.info("--- MEMORY ---")
     path = ensure_list(state, 'execution_path') + ["Memory"]
     mems = memory_manager.retrieve_relevant_memories(state.get('userInput',''))
     context = "\n".join([f"Memory: {m.page_content}" for m in mems]) if mems else "No memories"
-    return {"retrievedMemory": context, "execution_path": path, "status_update": "Memory retrieved"}
+    return {
+        "retrievedMemory": context, 
+        "execution_path": path,
+        **add_status_update("Memory", "Memory retrieved")
+    }
 
 def run_intent_agent(state: AgentState):
     log.info("--- INTENT ---")
@@ -407,7 +451,11 @@ def run_intent_agent(state: AgentState):
     prompt = f"Refine into clear objective.\n\nMemory: {state.get('retrievedMemory')}\n\nRequest: {state.get('userInput','')}\n\nCore Objective:"
     response = llm.invoke(prompt)
     core_obj = getattr(response, "content", "") or ""
-    return {"coreObjectivePrompt": core_obj, "execution_path": path, "status_update": "Objective clarified"}
+    return {
+        "coreObjectivePrompt": core_obj, 
+        "execution_path": path,
+        **add_status_update("Intent", "Objective clarified")
+    }
 
 def run_pm_agent(state: AgentState):
     log.info("--- PM ---")
@@ -424,14 +472,17 @@ def run_pm_agent(state: AgentState):
             "experiment_type": "word",
             "experiment_goal": state.get('coreObjectivePrompt', state.get('userInput',''))
         }
-        return {"pmPlan": fallback_plan, "execution_path": path, "rework_cycles": current_rework, "status_update": "Rework limit hit - manual review"}
+        return {
+            "pmPlan": fallback_plan, 
+            "execution_path": path, 
+            "rework_cycles": current_rework,
+            **add_status_update("PM", "Rework limit hit - manual review")
+        }
 
     # Normal behavior: increment rework count for this pass
     current_cycles = current_rework + 1
     path = ensure_list(state, 'execution_path') + ["PM"]
 
-    # (rest of your original PM prompt & parse flow, but ensure the output sets rework_cycles and max_loops)
-    # --- build full_context like before ---
     context_parts = [
         f"=== USER REQUEST ===\n{state.get('userInput', '')}",
         f"\n=== OBJECTIVE ===\n{state.get('coreObjectivePrompt', '')}",
@@ -481,8 +532,14 @@ Be concrete.
 
     # Attach loop control info
     plan['max_loops_initial'] = max_loops_val
-    plan['estimated_cost_usd'] = plan.get('estimated_cost_usd', plan.get('estimated_cost_usd', 0.0))
-    return {"pmPlan": plan, "execution_path": path, "rework_cycles": current_cycles, "max_loops": max_loops_val, "status_update": f"Plan created ({len(plan.get('plan_steps', []))} steps)"}
+    plan['estimated_cost_usd'] = plan.get('estimated_cost_usd', 0.0)
+    return {
+        "pmPlan": plan, 
+        "execution_path": path, 
+        "rework_cycles": current_cycles, 
+        "max_loops": max_loops_val,
+        **add_status_update("PM", f"Plan created ({len(plan.get('plan_steps', []))} steps)")
+    }
 
 def _extract_code_blocks(text: str, lang_hint: Optional[str]=None) -> List[str]:
     if lang_hint and "python" in (lang_hint or "").lower():
@@ -497,7 +554,12 @@ def run_experimenter_agent(state: AgentState):
     pm = state.get('pmPlan', {}) or {}
     
     if not pm.get('experiment_needed'):
-        return {"experimentCode": None, "experimentResults": None, "execution_path": path, "status_update": "No experiment needed"}
+        return {
+            "experimentCode": None, 
+            "experimentResults": None, 
+            "execution_path": path,
+            **add_status_update("Experimenter", "No experiment needed")
+        }
     
     exp_type = normalize_experiment_type(pm.get('experiment_type'), pm.get('experiment_goal',''))
     goal = pm.get('experiment_goal', 'No goal')
@@ -527,184 +589,161 @@ def run_experimenter_agent(state: AgentState):
 GOAL: {goal}
 
 CRITICAL REQUIREMENTS:
-
 1. ACTUAL WORKING CODE - Not templates, not documentation, not examples. REAL production code.
-
 2. FILE STRUCTURE - Indicate each file clearly:
    ### path/to/file.py
    ```python
    [Complete working code]
-   MUST INCLUDE:
-
-    Complete API clients with error handling, retries, rate limiting
-
-    Database schema with CREATE TABLE statements
-
-    Data processing with real transformation logic
-
-    Config management (.env handling)
-
-    requirements.txt with ALL dependencies
-
-    main.py entry point
-
-    Comprehensive README
-
-    CODE QUALITY:
-
-    Environment variables for secrets
-
-    Error handling and logging
-
-    Docstrings and comments
-
-    Real business logic based on request
+   ```
+
+MUST INCLUDE:
+- Complete API clients with error handling, retries, rate limiting
+- Database schema with CREATE TABLE statements
+- Data processing with real transformation logic
+- Config management (.env handling)
+- requirements.txt with ALL dependencies
+- main.py entry point
+- Comprehensive README
+
+CODE QUALITY:
+- Environment variables for secrets
+- Error handling and logging
+- Docstrings and comments
+- Real business logic based on request
+- RUNNABLE out of the box
+
+SPECIFIC TO REQUEST:
+- Use EXACT APIs mentioned (e.g., CricAPI, SportsRadar)
+- Implement SPECIFIC algorithms (e.g., batting avg, strike rate)
+- Create EXACT database tables needed
+- Process SPECIFIC data formats
+- NO placeholders like "# TODO"
+- NO dummy data - implement REAL logic
+- NO documentation-style code - PRODUCTION code only
+
+Format each file:
+### path/to/file.py
+```
+# Complete code here
+```
+
+Generate complete repository:"""
+
+        response = llm.invoke(repo_prompt)
+        llm_text = getattr(response, "content", "") or ""
+        
+        # Parse files from response
+        repo_files = {}
+        
+        # Extract with ### headers
+        file_pattern = r"###\s+([\w\/_\-\.]+)\s*\n```(?:\w+)?\s*\n(.*?)\n```"
+        matches = re.finditer(file_pattern, llm_text, re.DOTALL)
+        
+        for match in matches:
+            filepath = match.group(1).strip()
+            content = match.group(2).strip()
+            repo_files[filepath] = content
+        
+        # Fallback: extract code blocks
+        if not repo_files:
+            code_blocks = re.findall(r"```(?:python|sql)?\s*\n(.*?)\n```", llm_text, re.DOTALL)
+            if code_blocks:
+                for i, block in enumerate(code_blocks):
+                    if len(block) > 50:  # Skip tiny blocks
+                        repo_files[f"module_{i}.py"] = block
+        
+        # Add README if missing
+        if not any('README' in f.upper() for f in repo_files):
+            repo_files["README.md"] = f"""# Generated Application
 
-    RUNNABLE out of the box
+## Overview
+{goal}
 
-    SPECIFIC TO REQUEST:
+## Files
+{chr(10).join(f'- {f}' for f in sorted(repo_files.keys()))}
 
-    Use EXACT APIs mentioned (e.g., CricAPI, SportsRadar)
+## Setup
+1. `pip install -r requirements.txt`
+2. Copy `.env.example` to `.env` and configure
+3. Run: `python main.py`
+"""
 
-    Implement SPECIFIC algorithms (e.g., batting avg, strike rate)
+        # Add requirements.txt if missing
+        if "requirements.txt" not in repo_files:
+            all_code = " ".join(repo_files.values()).lower()
+            deps = []
+            if 'requests' in all_code: deps.append('requests')
+            if 'pandas' in all_code: deps.append('pandas')
+            if 'numpy' in all_code: deps.append('numpy')
+            if 'sqlalchemy' in all_code: deps.append('sqlalchemy')
+            if 'postgresql' in all_code or 'psycopg2' in all_code: deps.append('psycopg2-binary')
+            if 'flask' in all_code: deps.append('flask')
+            if 'fastapi' in all_code: 
+                deps.append('fastapi')
+                deps.append('uvicorn')
+            if 'dotenv' in all_code: deps.append('python-dotenv')
+            
+            repo_files["requirements.txt"] = "\n".join(deps) if deps else "# Dependencies"
+
+        # Add .env.example if missing
+        if ".env.example" not in repo_files:
+            repo_files[".env.example"] = """# Configuration
+API_KEY=your_key_here
+DATABASE_URL=postgresql://user:pass@localhost/db
+DEBUG=False
+"""
 
-    Create EXACT database tables needed
+        # Add main.py if missing
+        if not any('main.py' in f for f in repo_files):
+            repo_files["main.py"] = """#!/usr/bin/env python3
+import os
+from dotenv import load_dotenv
 
-    Process SPECIFIC data formats
+load_dotenv()
 
-    NO placeholders like "# TODO"
-    NO dummy data - implement REAL logic
-    NO documentation-style code - PRODUCTION code only
+def main():
+    print("Application starting...")
+    # Add your logic here
+    pass
 
-    Format each file:
+if __name__ == "__main__":
+    main()
+"""
 
-    path/to/file.py
-    # Complete code here
-    Generate complete repository:"""
+        # Build zip
+        zip_path = build_repo_zip(repo_files, repo_name="generated_app", out_dir=OUT_DIR)
+        
+        results = {
+            "success": True,
+            "paths": {"repo_zip": sanitize_path(zip_path)},
+            "files_created": len(repo_files),
+            "context_used": len(full_context)
+        }
+        
+        return {
+            "experimentCode": None,
+            "experimentResults": results,
+            "execution_path": path,
+            **add_status_update("Experimenter", f"Repository created ({len(repo_files)} files)")
+        }
 
     # OTHER ARTIFACT TYPES
     enhanced_prompt = f"""Create HIGH-QUALITY {exp_type} artifact.
-    {full_context}
-
-    GOAL: {goal}
-
-    REQUIREMENTS:
-
-    Use ALL specific details from request
-
-    PRODUCTION-READY, COMPLETE content (NO placeholders)
-
-    ACTUAL data, REALISTIC examples, WORKING code
-
-    For notebooks: markdown + executable code + visualizations
-
-    For scripts: error handling + docs + real logic
-
-    For documents: substantive detailed content
-
-    Generate complete content for '{exp_type}' with proper code fences."""
 
+{full_context}
 
-    response = llm.invoke(enhanced_prompt)
-    llm_text = getattr(response, "content", "") or ""
-    
-    # Parse files from response
-    repo_files = {}
-    
-    # Extract with ### headers
-    file_pattern = r"###\s+([\w\/_\-\.]+)\s*\n```(?:\w+)?\s*\n(.*?)\n```"
-    matches = re.finditer(file_pattern, llm_text, re.DOTALL)
-    
-    for match in matches:
-        filepath = match.group(1).strip()
-        content = match.group(2).strip()
-        repo_files[filepath] = content
-    
-    # Fallback: extract code blocks
-    if not repo_files:
-        code_blocks = re.findall(r"```(?:python|sql)?\s*\n(.*?)\n```", llm_text, re.DOTALL)
-        if code_blocks:
-            for i, block in enumerate(code_blocks):
-                if len(block) > 50:  # Skip tiny blocks
-                    repo_files[f"module_{i}.py"] = block
-    
-    # Add README if missing
-    if not any('README' in f.upper() for f in repo_files):
-        repo_files["README.md"] = f"""# Generated Application
-                                        Overview
-                                        {goal}
-
-                                        Files
-                                        {chr(10).join(f'- {f}' for f in sorted(repo_files.keys()))}
-
-                                        Setup
-                                        pip install -r requirements.txt
-
-                                        Copy .env.example to .env and configure
-
-                                        Run: python main.py
-                                        """
-
-     # Add requirements.txt
-    if "requirements.txt" not in repo_files:
-         all_code = " ".join(repo_files.values()).lower()
-         deps = []
-         if 'requests' in all_code: deps.append('requests')
-         if 'pandas' in all_code: deps.append('pandas')
-         if 'numpy' in all_code: deps.append('numpy')
-         if 'sqlalchemy' in all_code: deps.append('sqlalchemy')
-         if 'postgresql' in all_code or 'psycopg2' in all_code: deps.append('psycopg2-binary')
-         if 'flask' in all_code: deps.append('flask')
-         if 'fastapi' in all_code: 
-             deps.append('fastapi')
-             deps.append('uvicorn')
-         if 'dotenv' in all_code: deps.append('python-dotenv')
-
-         repo_files["requirements.txt"] = "\n".join(deps) if deps else "# Dependencies"
-
-     # Add .env.example
-    if ".env.example" not in repo_files:
-         repo_files[".env.example"] = """# Configuration
-    API_KEY=your_key_here
-    DATABASE_URL=postgresql://user:pass@localhost/db
-    DEBUG=False
-    """
-
-    # Add main.py if missing
-    if not any('main.py' in f for f in repo_files):
-            repo_files["main.py"] = """#!/usr/bin/env python3
-    import os
-    from dotenv import load_dotenv
-
-    load_dotenv()
-
-    def main():
-    print("Application starting...")
-    # Add your logic here
-    pass
+GOAL: {goal}
 
-    if name == "main":
-    main()
-    """
+REQUIREMENTS:
+- Use ALL specific details from request
+- PRODUCTION-READY, COMPLETE content (NO placeholders)
+- ACTUAL data, REALISTIC examples, WORKING code
+- For notebooks: markdown + executable code + visualizations
+- For scripts: error handling + docs + real logic
+- For documents: substantive detailed content
 
-    # Build zip
-    zip_path = build_repo_zip(repo_files, repo_name="generated_app", out_dir=OUT_DIR)
-    
-    results = {
-        "success": True,
-        "paths": {"repo_zip": sanitize_path(zip_path)},
-        "files_created": len(repo_files),
-        "context_used": len(full_context)
-    }
-    
-    return {
-        "experimentCode": None,
-        "experimentResults": results,
-        "execution_path": path,
-        "status_update": f"Repository created ({len(repo_files)} files)"
-    }
+Generate complete content for '{exp_type}' with proper code fences."""
 
-    
     response = llm.invoke(enhanced_prompt)
     llm_text = getattr(response, "content", "") or ""
     results = {"success": False, "paths": {}, "stderr": "", "stdout": "", "context_used": len(full_context)}
@@ -713,22 +752,42 @@ CRITICAL REQUIREMENTS:
         if exp_type == 'notebook':
             nb_path = write_notebook_from_text(llm_text, out_dir=OUT_DIR)
             results.update({"success": True, "paths": {"notebook": sanitize_path(nb_path)}})
-            return {"experimentCode": None, "experimentResults": results, "execution_path": path, "status_update": "Notebook created"}
+            return {
+                "experimentCode": None, 
+                "experimentResults": results, 
+                "execution_path": path,
+                **add_status_update("Experimenter", "Notebook created")
+            }
         
         elif exp_type == 'excel':
             excel_path = write_excel_from_tables(llm_text, out_dir=OUT_DIR)
             results.update({"success": True, "paths": {"excel": sanitize_path(excel_path)}})
-            return {"experimentCode": None, "experimentResults": results, "execution_path": path, "status_update": "Excel created"}
+            return {
+                "experimentCode": None, 
+                "experimentResults": results, 
+                "execution_path": path,
+                **add_status_update("Experimenter", "Excel created")
+            }
         
         elif exp_type == 'word':
             docx_path = write_docx_from_text(llm_text, out_dir=OUT_DIR)
             results.update({"success": True, "paths": {"docx": sanitize_path(docx_path)}})
-            return {"experimentCode": None, "experimentResults": results, "execution_path": path, "status_update": "DOCX created"}
+            return {
+                "experimentCode": None, 
+                "experimentResults": results, 
+                "execution_path": path,
+                **add_status_update("Experimenter", "Word document created")
+            }
         
         elif exp_type == 'pdf':
             pdf_path = write_pdf_from_text(llm_text, out_dir=OUT_DIR)
             results.update({"success": True, "paths": {"pdf": sanitize_path(pdf_path)}})
-            return {"experimentCode": None, "experimentResults": results, "execution_path": path, "status_update": "PDF created"}
+            return {
+                "experimentCode": None, 
+                "experimentResults": results, 
+                "execution_path": path,
+                **add_status_update("Experimenter", "PDF created")
+            }
         
         elif exp_type == 'script':
             lang_hint = pm.get('experiment_language') or "python"
@@ -750,17 +809,33 @@ CRITICAL REQUIREMENTS:
                 "stdout": exec_results.get("stdout",""),
                 "stderr": exec_results.get("stderr","")
             })
-            return {"experimentCode": code_text, "experimentResults": results, "execution_path": path, "status_update": "Script created"}
+            return {
+                "experimentCode": code_text, 
+                "experimentResults": results, 
+                "execution_path": path,
+                **add_status_update("Experimenter", "Script created")
+            }
     
         else:
             fallback = write_docx_from_text(llm_text, out_dir=OUT_DIR)
             results.update({"success": True, "paths": {"docx": sanitize_path(fallback)}})
-            return {"experimentCode": None, "experimentResults": results, "execution_path": path, "status_update": "Document created"}
+            return {
+                "experimentCode": None, 
+                "experimentResults": results, 
+                "execution_path": path,
+                **add_status_update("Experimenter", "Document created")
+            }
         
     except Exception as e:
         log.error(f"Experimenter failed: {e}")
         results.update({"success": False, "stderr": str(e)})
-        return {"experimentCode": None, "experimentResults": results, "execution_path": path, "status_update": "Error"}
+        return {
+            "experimentCode": None, 
+            "experimentResults": results, 
+            "execution_path": path,
+            **add_status_update("Experimenter", f"Error: {str(e)}")
+        }
+
 def run_synthesis_agent(state: AgentState):
     log.info("--- SYNTHESIS ---")
     _state = state or {}
@@ -799,23 +874,17 @@ def run_synthesis_agent(state: AgentState):
     full_context = "\n".join(synthesis_context)
 
     synthesis_prompt = f"""Create FINAL RESPONSE after executing user's request.
-    {full_context}
-
-    Create comprehensive response that:
-
-    Directly addresses original request
 
-    Explains what was accomplished and HOW
-
-    References specific artifacts and explains PURPOSE
-
-    Provides context on how to USE deliverables
-
-    Highlights KEY INSIGHTS
-
-    Suggests NEXT STEPS if relevant
+{full_context}
 
-    Be SPECIFIC about what was created."""
+Create comprehensive response that:
+- Directly addresses original request
+- Explains what was accomplished and HOW
+- References specific artifacts and explains PURPOSE
+- Provides context on how to USE deliverables
+- Highlights KEY INSIGHTS
+- Suggests NEXT STEPS if relevant
+- Be SPECIFIC about what was created."""
 
     response = llm.invoke(synthesis_prompt)
     final_text = getattr(response, "content", "") or ""
@@ -823,7 +892,11 @@ def run_synthesis_agent(state: AgentState):
     if artifact_message:
         final_text = final_text + "\n\n---\n" + artifact_message
 
-    return {"draftResponse": final_text, "execution_path": path, "status_update": "Response synthesized"}
+    return {
+        "draftResponse": final_text, 
+        "execution_path": path,
+        **add_status_update("Synthesis", "Response synthesized")
+    }
 
 def run_qa_agent(state: AgentState):
     log.info("--- QA ---")
@@ -839,40 +912,47 @@ def run_qa_agent(state: AgentState):
         qa_context.append(f"\n=== ARTIFACTS ===\n{json.dumps(state.get('experimentResults', {}).get('paths', {}), indent=2)}")
 
     prompt = f"""You are a QA reviewer. Review the draft response against the user's objective.
-    {chr(10).join(qa_context)}
 
-    Review Instructions:
+{chr(10).join(qa_context)}
 
-    Does the draft and its artifacts COMPLETELY satisfy ALL parts of the user's request?
+Review Instructions:
+- Does the draft and its artifacts COMPLETELY satisfy ALL parts of the user's request?
+- Is the quality of the work high?
+- If this is a re-submission (rework cycle > 1), has the previous feedback been successfully addressed?
 
-    Is the quality of the work high?
+Response Format (required JSON or a single word 'APPROVED'):
 
-    If this is a re-submission (rework cycle > 1), has the previous feedback been successfully addressed?
+Either return EXACTLY the single word:
+  APPROVED
 
-    Response Format (required JSON or a single word 'APPROVED'):
-
-    Either return EXACTLY the single word:
-      APPROVED
-
-    Or return JSON like:
-    {{
-      "approved": false,
-      "feedback": "Specific, actionable items to fix (bullet list or numbered).",
-      "required_changes": ["..."]
-    }}
-    """
+Or return JSON like:
+{{
+  "approved": false,
+  "feedback": "Specific, actionable items to fix (bullet list or numbered).",
+  "required_changes": ["..."]
+}}
+"""
 
     try:
         response = llm.invoke(prompt)
         content = getattr(response, "content", "") or ""
     except Exception as e:
         log.exception("QA LLM call failed: %s", e)
-        # Fail-safe: mark as not approved with conservative feedback
-        return {"approved": False, "qaFeedback": "QA LLM failed; manual review required.", "execution_path": path, "status_update": "QA failed"}
+        return {
+            "approved": False, 
+            "qaFeedback": "QA LLM failed; manual review required.", 
+            "execution_path": path,
+            **add_status_update("QA", "QA failed")
+        }
 
     # If LLM returned APPROVED word, treat as approved
     if "APPROVED" in content.strip().upper() and len(content.strip()) <= 20:
-        return {"approved": True, "qaFeedback": None, "execution_path": path, "status_update": "Approved"}
+        return {
+            "approved": True, 
+            "qaFeedback": None, 
+            "execution_path": path,
+            **add_status_update("QA", "Approved")
+        }
 
     # Else try JSON parse
     parsed = parse_json_from_llm(content)
@@ -884,10 +964,21 @@ def run_qa_agent(state: AgentState):
             feedback = "\n".join([str(x) for x in feedback])
         elif not isinstance(feedback, str):
             feedback = str(feedback)
-        return {"approved": approved, "qaFeedback": feedback if not approved else None, "execution_path": path, "status_update": "QA completed"}
+        return {
+            "approved": approved, 
+            "qaFeedback": feedback if not approved else None, 
+            "execution_path": path,
+            **add_status_update("QA", "QA completed")
+        }
+    
     # Fallback: return raw text as feedback (not approved)
     safe_feedback = content.strip()[:2000] or "QA produced no actionable output."
-    return {"approved": False, "qaFeedback": safe_feedback, "execution_path": path, "status_update": "QA needs rework"}
+    return {
+        "approved": False, 
+        "qaFeedback": safe_feedback, 
+        "execution_path": path,
+        **add_status_update("QA", "QA needs rework")
+    }
 
 def run_archivist_agent(state: AgentState):
     log.info("--- ARCHIVIST ---")
@@ -897,7 +988,10 @@ def run_archivist_agent(state: AgentState):
     response = llm.invoke(summary_prompt)
     memory_manager.add_to_memory(getattr(response,"content",""), {"objective": state.get('coreObjectivePrompt')})
 
-    return {"execution_path": path, "status_update": "Saved to memory"}
+    return {
+        "execution_path": path,
+        **add_status_update("Archivist", "Saved to memory")
+    }
 
 def run_disclaimer_agent(state: AgentState):
     log.warning("--- DISCLAIMER ---")
@@ -907,7 +1001,11 @@ def run_disclaimer_agent(state: AgentState):
     disclaimer = f"**DISCLAIMER: {reason} Draft may be incomplete.**\n\n---\n\n"
     final_response = disclaimer + state.get('draftResponse', "No response")
 
-    return {"draftResponse": final_response, "execution_path": path, "status_update": reason}
+    return {
+        "draftResponse": final_response, 
+        "execution_path": path,
+        **add_status_update("Disclaimer", reason)
+    }
 
 def should_continue(state: AgentState):
     # Budget check first
@@ -929,7 +1027,6 @@ def should_continue(state: AgentState):
     # Default: return pm_agent so planner will create next plan
     return "pm_agent"
 
-
 def should_run_experiment(state: AgentState):
     pm = state.get('pmPlan', {}) or {}
     return "experimenter_agent" if pm.get('experiment_needed') else "synthesis_agent"
@@ -967,11 +1064,9 @@ main_workflow.add_edge("disclaimer_agent", END)
 
 main_workflow.add_conditional_edges("pm_agent", should_run_experiment)
 main_workflow.add_conditional_edges("qa_agent", should_continue, {
-"archivist_agent": "archivist_agent",
-"pm_agent": "pm_agent",
-"disclaimer_agent": "disclaimer_agent"
+    "archivist_agent": "archivist_agent",
+    "pm_agent": "pm_agent",
+    "disclaimer_agent": "disclaimer_agent"
 })
 
-main_app = main_workflow.compile()
-
-
+main_app = main_workflow.compile()