"""LangGraph orchestration integration for CyberBrain26 Phase 6. Implements multi-agent orchestration using LangGraph patterns (subagents, supervisor-worker hierarchy) with shared AgentState for state propagation. Integrates with existing PtahBridge / AgentRegistry / TaskDispatcher security chain. """ from __future__ import annotations import asyncio import logging from typing import TYPE_CHECKING, Any if TYPE_CHECKING: from langgraph.graph import StateGraph, END from langgraph.graph.state import CompiledStateGraph from cyberbrain26.agents.orchestrator import BaseAgent, AgentResult from cyberbrain26.control_plane import AsyncEventBus, MemoryTaskQueue from cyberbrain26.security_plane.identity import AgentIdentity from cyberbrain26.security_plane.permission import PermissionGateway logger = logging.getLogger(__name__) # Lazily imported because langgraph is an optional capability. _langgraph_available: bool | None = None def _ensure_langgraph() -> tuple[type, object]: """Import langgraph lazily; raise a clear error if missing.""" global _langgraph_available if _langgraph_available is None: try: from langgraph.graph import StateGraph, END # noqa: PLC0415 from langgraph.graph.state import CompiledStateGraph # noqa: PLC0415 _langgraph_available = (StateGraph, END, CompiledStateGraph) except ImportError as exc: _langgraph_available = False raise RuntimeError( "LangGraph not installed; orchestration capability unavailable." ) from exc if _langgraph_available is False: raise RuntimeError("LangGraph not installed; orchestration capability unavailable.") return _langgraph_available class AgentState(dict[str, Any]): """Shared state propagated between supervisor and worker agents. Carries task data, intermediate results, and routing decisions through the LangGraph orchestration loop. """ def __init__(self, **kwargs: Any) -> None: super().__init__( { "task_id": None, "goal": None, "current_agent": None, "results": [], "status": "pending", "error": None, }, **kwargs, ) class LangGraphOrchestrator: """Orchestrates agent execution via LangGraph state graph. Builds a supervisor → worker hierarchy where: 1. Supervisor analyzes the goal and assigns tasks to workers 2. Workers execute specialized tasks (coder, evaluator, researcher, etc.) 3. Results are aggregated and routed back through supervisor Integrates with CyberBrain26 security plane: every worker invocation passes through AgentIdentity + PermissionGateway before execution. """ def __init__( self, bus: AsyncEventBus, queue: MemoryTaskQueue, permission_gateway: PermissionGateway, agents: dict[str, BaseAgent], ) -> None: self.bus = bus self.queue = queue self.permission_gateway = permission_gateway self.agents = agents self._graph: CompiledStateGraph | None = None self._build_graph() def _build_graph(self) -> None: """Build LangGraph state graph with supervisor and worker nodes.""" StateGraph, END, _ = _ensure_langgraph() async def supervisor_node(state: AgentState) -> AgentState: """Supervisor analyzes state and routes to appropriate worker.""" goal = state.get("goal", "") task_id = state.get("task_id", "") # Security check: verify supervisor may dispatch identity = AgentIdentity(agent_id="langgraph-supervisor", role="supervisor") if not self.permission_gateway.evaluate(identity, "dispatch", goal).decision.value == "allow": state["status"] = "blocked" state["error"] = "supervisor denied dispatch permission" return state # Simple routing logic based on goal keywords content = goal.lower() if any(w in content for w in ["code", "function", "implement", "write"]): state["current_agent"] = "coder" elif any(w in content for w in ["research", "search", "web", "find"]): state["current_agent"] = "researcher" elif any(w in content for w in ["evaluate", "assess", "review", "check"]): state["current_agent"] = "evaluator" else: state["current_agent"] = "planner" return state async def worker_node(state: AgentState) -> AgentState: """Worker agent executes assigned task through security gate.""" agent_role = state.get("current_agent", "planner") agent = self.agents.get(agent_role) if agent is None: state["status"] = "failed" state["error"] = f"unknown agent role: {agent_role}" return state # Security gate: check agent permission identity = AgentIdentity(agent_id=f"langgraph-{agent_role}", role=agent_role) decision = self.permission_gateway.evaluate(identity, "execute", state.get("goal", "")) if decision.decision.value != "allow": state["status"] = "blocked" state["error"] = f"agent {agent_role} denied execution permission" return state # Execute via agent's execute method task_data = { "task_id": state.get("task_id"), "goal": state.get("goal"), "max_steps": 3, } try: result = await agent.execute(task_data) if result: state["results"].append( { "agent": agent.name, "role": agent_role, "success": result.success, "data": result.data, "error": result.error, } ) state["status"] = ( "completed" if result.success else "failed" ) state["error"] = result.error except Exception as exc: # noqa: BLE001 - worker failure captured state["status"] = "failed" state["error"] = str(exc) return state def router(state: AgentState) -> str: """Route to next node based on state status.""" status = state.get("status", "pending") if status in ("completed", "failed", "blocked"): return END return "worker" graph = StateGraph(AgentState) graph.add_node("supervisor", supervisor_node) graph.add_node("worker", worker_node) graph.add_edge("supervisor", "worker") graph.add_conditional_edges("worker", router) graph.set_entry_point("supervisor") self._graph = graph.compile() async def run(self, task_id: str, goal: str) -> AgentState: """Execute orchestrated task through LangGraph graph.""" if self._graph is None: raise RuntimeError("LangGraph graph not built") initial_state = AgentState( task_id=task_id, goal=goal, current_agent=None, results=[], status="pending", ) final_state = await self._graph.ainvoke(initial_state) return final_state async def shutdown(self) -> None: """Cleanup orchestrator resources.""" logger.debug("LangGraphOrchestrator shutdown") __all__ = ["LangGraphOrchestrator", "AgentState"]