core/agents/king_agent.py

"""
King Agent — Sovereign Executive Orchestrator

The King Agent is the counter-part to the Queen Agent. While the Queen
architectures the solution (blueprints), the King executes it by 
orchestrating specialty agents and enforcing sovereign governance.
"""

import logging
import asyncio
from typing import Dict, Any, List, Optional
from core.agents.queen_agent import QueenAgent
from core.blueprint_healer import BlueprintHealer
from core.models import User, AgentTriggerMode, AgentEvolutionTrace
from core.database import SessionLocal
from core.atom_meta_agent import AtomMetaAgent
from tools.canvas_tool import present_markdown, update_canvas

logger = logging.getLogger(__name__)

class KingAgent(AtomMetaAgent):
    """
    The King Agent takes a blueprint and oversees its realization.
    It manages the lifecycle of sub-tasks and ensures each node in the 
    blueprint is executed by the most capable available resource.
    """
    
    def __init__(self, workspace_id: str = "default", tenant_id: Optional[str] = None, user: Optional[User] = None):
        super().__init__(workspace_id, tenant_id, user)
        # LLM initialized in AtomMetaAgent via ServiceFactory
        self.healer = BlueprintHealer(None, self.llm) # db will be injected in session

    async def execute_blueprint(
        self, 
        blueprint: Dict[str, Any], 
        context: Optional[Dict[str, Any]] = None,
        canvas_context: Optional[Dict[str, str]] = None
    ) -> Dict[str, Any]:
        """
        Executes a Queen's blueprint with sovereign self-healing.
        """
        logger.info(f"King: Executing blueprint: {blueprint.get('architecture_name')}")
        context = context or {}
        nodes = blueprint.get("nodes", [])
        
        executed_nodes = {}
        pending_nodes = nodes.copy()
        
        # Prepare initial Canvas visualization
        canvas_id = None
        node_statuses = {n["id"]: "pending" for n in nodes}
        
        if context.get("user_id") and context.get("tenant_id"):
            mermaid = self.queen.generate_mermaid(blueprint, node_statuses)
            canvas_res = await present_markdown(
                tenant_id=context["tenant_id"],
                user_id=context["user_id"],
                title=f"Execution Plan: {blueprint.get('architecture_name')}",
                content=f"```mermaid\n{mermaid}\n```",
                agent_id=getattr(self, "agent_id", None)
            )
            canvas_id = canvas_res.get("canvas_id")

        while pending_nodes and retry_count < max_retries:
            ready_nodes = [
                n for n in pending_nodes 
                if all(dep in executed_nodes for dep in n.get("dependencies", []))
            ]
            
            if not ready_nodes and pending_nodes:
                logger.error("King: Stalled execution or circular dependency.")
                break
                
            for node in ready_nodes:
                logger.info(f"King: Processing node: {node['name']} ({node['type']})")
                
                # Update status to in_progress on Canvas
                if canvas_id:
                    node_statuses[node["id"]] = "in_progress"
                    new_mermaid = self.queen.generate_mermaid(blueprint, node_statuses)
                    await update_canvas(
                        tenant_id=context["tenant_id"],
                        user_id=context["user_id"],
                        canvas_id=canvas_id,
                        updates={"content": f"```mermaid\n{new_mermaid}\n```"}
                    )

                try:
                    # Execute the node
                    node_result = await self._execute_node(node, context, canvas_context)
                    
                    if isinstance(node_result, dict) and "error" in node_result:
                        raise ValueError(node_result["error"])

                    executed_nodes[node["id"]] = node_result
                    pending_nodes.remove(node)
                    node_statuses[node["id"]] = "completed"
                    
                    results.append({
                        "node_id": node["id"],
                        "node_name": node["name"],
                        "status": "completed",
                        "result": node_result
                    })
                    
                    # Update Canvas with completion
                    if canvas_id:
                        new_mermaid = self.queen.generate_mermaid(blueprint, node_statuses)
                        await update_canvas(
                            tenant_id=context["tenant_id"],
                            user_id=context["user_id"],
                            canvas_id=canvas_id,
                            updates={"content": f"```mermaid\n{new_mermaid}\n```"}
                        )

                except Exception as e:
                    logger.warning(f"King: Node failure detected in {node['id']}: {e}")
                    node_statuses[node["id"]] = "failed"
                    
                    # Update Canvas with failure
                    if canvas_id:
                        new_mermaid = self.queen.generate_mermaid(blueprint, node_statuses)
                        await update_canvas(
                            tenant_id=context["tenant_id"],
                            user_id=context["user_id"],
                            canvas_id=canvas_id,
                            updates={"content": f"```mermaid\n{new_mermaid}\n```"}
                        )

                    # TRIGGER SELF-HEALING
                    logger.info("King: Activating Blueprint Healer...")
                    healed_blueprint = await self.healer.heal_blueprint(
                        blueprint=blueprint,
                        failed_node_id=node["id"],
                        error_message=str(e),
                        tenant_id=context.get("tenant_id", "default")
                    )
                    
                    if healed_blueprint.get("status") == "healed":
                        logger.info("King: Blueprint healed. Record learning trace.")
                        
                        # RECORD LEARNING TRACE
                        try:
                            directive = await self.healer.summarize_healing_as_directive(
                                failed_node=node,
                                healed_nodes=healed_blueprint.get("nodes", []),
                                error=str(e),
                                tenant_id=context.get("tenant_id", "default")
                            )
                            
                            with SessionLocal() as db:
                                trace = AgentEvolutionTrace(
                                    tenant_id=context.get("tenant_id", "default"),
                                    agent_id=getattr(self, "agent_id", "king_agent"),
                                    evolution_type="performance_based",
                                    task_log=f"Node failure: {node['id']}\nError: {e}",
                                    evolving_requirements=directive,
                                    model_patch=json.dumps(healed_blueprint.get("nodes", [])),
                                    benchmark_passed=True, 
                                    benchmark_score=1.0,
                                    is_high_quality=True
                                )
                                db.add(trace)
                                db.commit()
                                logger.info(f"King: Recorded evolution trace with directive: {directive}")
                        except Exception as tracing_err:
                            logger.error(f"King: Failed to record learning trace: {tracing_err}")

                        logger.info("King: Restarting execution with patched architecture.")
                        blueprint = healed_blueprint
                        nodes = blueprint.get("nodes", [])
                        # Re-sync node statuses
                        for n in nodes:
                            if n["id"] not in node_statuses:
                                node_statuses[n["id"]] = "pending"
                        
                        pending_nodes = [n for n in nodes if n["id"] not in executed_nodes]
                        retry_count += 1
                        
                        # Update Canvas with new blueprint
                        if canvas_id:
                            new_mermaid = self.queen.generate_mermaid(blueprint, node_statuses)
                            await update_canvas(
                                tenant_id=context["tenant_id"],
                                user_id=context["user_id"],
                                canvas_id=canvas_id,
                                updates={
                                    "title": f"Healed Plan ({retry_count}): {blueprint.get('architecture_name')}",
                                    "content": f"```mermaid\n{new_mermaid}\n```"
                                }
                            )
                        break 
                    else:
                        logger.error("King: Healing failed. Abandoning execution.")
                        return {
                            "status": "failed",
                            "error": str(e),
                            "partial_results": results
                        }

        return {
            "status": "success",
            "blueprint_id": blueprint.get("blueprint_id"),
            "execution_results": results,
            "final_summary": f"Blueprint '{blueprint.get('architecture_name')}' executed successfully with {retry_count} heal events."
        }

    async def _execute_node(self, node: Dict[str, Any], context: Dict, canvas_context: Optional[Dict]) -> Any:
        """
        Executes a single node in the blueprint.
        """
        node_id = node.get("id")
        node_type = node.get("type", "agent")
        capability = node.get("capability_required")
        
        # If it's an agent node, delegate to a specialized agent
        if node_type == "agent":
            # Map capability to agent name if needed
            agent_name = self._map_capability_to_agent(capability)
            logger.info(f"King: Delegating '{node['name']}' to {agent_name}")
            
            # Use parent's delegation logic
            return await self._execute_delegation(
                agent_name=agent_name,
                task=f"Objective: {node['name']}. Requirements: {capability}",
                context=context
            )
        
        # If it's a direct skill/tool call
        elif node_type == "skill":
            logger.info(f"King: Executing skill: {node['name']}")
            # Use parent's tool execution (includes governance)
            return await self._execute_tool_with_governance(
                tool_name=capability,
                args=node.get("params", {}),
                context=context,
                step_callback=None
            )
            
        return {"status": "skipped", "reason": f"Unknown node type: {node_type}"}

    def _map_capability_to_agent(self, capability: str) -> str:
        """Helper to route capabilities to specialized agents."""
        mapping = {
            "reconciliation": "accounting",
            "lead_scoring": "sales",
            "inventory_check": "logistics",
            "campaign_analysis": "marketing",
            "b2b_extract_po": "purchasing"
        }
        return mapping.get(capability, "general") # Default to Meta (Self) or General