1"""
2Multi-Agent Coordination Utilities
3 
4Provides reusable building blocks for multi-agent coordination patterns:
5supervisor/orchestrator, peer-to-peer handoffs, consensus mechanisms,
6and failure handling with circuit breakers.
7 
8Use when: building multi-agent systems that need structured communication,
9task delegation, consensus voting, or fault-tolerant agent coordination.
10 
11Designed for composability — import individual classes or use the
12``if __name__ == "__main__"`` demo to see all patterns in action.
13"""
14 
15from typing import Dict, List, Any, Optional
16from dataclasses import dataclass, field
17from enum import Enum
18import time
19import uuid
20 
21__all__ = [
22    "MessageType",
23    "AgentMessage",
24    "AgentCommunication",
25    "SupervisorAgent",
26    "HandoffProtocol",
27    "ConsensusManager",
28    "AgentFailureHandler",
29]
30 
31 
32class MessageType(Enum):
33    """Types of messages exchanged between agents."""
34 
35    REQUEST = "request"
36    RESPONSE = "response"
37    HANDOVER = "handover"
38    FEEDBACK = "feedback"
39    ALERT = "alert"
40 
41 
42@dataclass
43class AgentMessage:
44    """Message exchanged between agents.
45 
46    Use when: agents need a structured envelope for inter-agent communication
47    that carries sender/receiver identity, type, priority, and payload.
48    """
49 
50    sender: str
51    receiver: str
52    message_type: MessageType
53    content: Dict[str, Any]
54    timestamp: float = field(default_factory=time.time)
55    message_id: str = field(default_factory=lambda: str(uuid.uuid4()))
56    requires_response: bool = False
57    priority: int = 0  # 0 = normal, higher = more urgent
58 
59 
60class AgentCommunication:
61    """Communication channel for multi-agent systems.
62 
63    Use when: multiple agents need an in-process message bus for sending,
64    receiving, and broadcasting messages with history tracking.
65    """
66 
67    def __init__(self) -> None:
68        self.inbox: Dict[str, List[AgentMessage]] = {}
69        self.outbox: List[AgentMessage] = []
70        self.message_history: List[AgentMessage] = []
71 
72    def send(self, message: AgentMessage) -> None:
73        """Send a message to an agent."""
74        if message.receiver not in self.inbox:
75            self.inbox[message.receiver] = []
76        self.inbox[message.receiver].append(message)
77        self.outbox.append(message)
78        self.message_history.append(message)
79 
80    def receive(self, agent_id: str) -> List[AgentMessage]:
81        """Receive all messages for an agent, clearing its inbox."""
82        messages = self.inbox.get(agent_id, [])
83        self.inbox[agent_id] = []
84        return messages
85 
86    def broadcast(
87        self,
88        sender: str,
89        message_type: MessageType,
90        content: Dict[str, Any],
91        receivers: List[str],
92    ) -> None:
93        """Broadcast a message to multiple agents."""
94        for receiver in receivers:
95            self.send(
96                AgentMessage(
97                    sender=sender,
98                    receiver=receiver,
99                    message_type=message_type,
100                    content=content,
101                )
102            )
103 
104 
105# ---------------------------------------------------------------------------
106# Supervisor Pattern
107# ---------------------------------------------------------------------------
108 
109 
110class SupervisorAgent:
111    """Central supervisor agent that coordinates worker agents.
112 
113    Use when: tasks have clear decomposition and a single coordinator should
114    delegate subtasks, track worker status, and aggregate results.
115    """
116 
117    def __init__(self, name: str, communication: AgentCommunication) -> None:
118        self.name = name
119        self.communication = communication
120        self.workers: Dict[str, Dict[str, Any]] = {}
121        self.task_queue: List[Dict[str, Any]] = []
122        self.completed_tasks: List[Dict[str, Any]] = []
123        self.current_state: Dict[str, Any] = {}
124 
125    def register_worker(self, worker_id: str, capabilities: List[str]) -> None:
126        """Register a worker agent with the supervisor."""
127        self.workers[worker_id] = {
128            "capabilities": capabilities,
129            "status": "available",
130            "current_task": None,
131            "metrics": {"tasks_completed": 0, "avg_response_time": 0.0},
132        }
133 
134    def decompose_task(self, task: Dict[str, Any]) -> List[Dict[str, Any]]:
135        """Decompose a task into subtasks.
136 
137        Use when: a high-level task needs to be broken into assignable units.
138        In production, replace the rule-based logic with LLM-driven planning.
139        """
140        subtasks: List[Dict[str, Any]] = []
141        task_type = task.get("type", "general")
142 
143        if task_type == "research":
144            subtasks = [
145                {"type": "search", "description": "Gather information"},
146                {"type": "analyze", "description": "Analyze findings"},
147                {"type": "synthesize", "description": "Synthesize results"},
148            ]
149        elif task_type == "create":
150            subtasks = [
151                {"type": "plan", "description": "Create plan"},
152                {"type": "draft", "description": "Draft content"},
153                {"type": "review", "description": "Review and refine"},
154            ]
155        else:
156            subtasks = [
157                {
158                    "type": "execute",
159                    "description": task.get("description", "Execute task"),
160                }
161            ]
162 
163        for subtask in subtasks:
164            subtask["parent_task"] = task.get("id")
165            subtask["priority"] = task.get("priority", 0)
166 
167        return subtasks
168 
169    def assign_task(self, subtask: Dict[str, Any], worker_id: str) -> None:
170        """Assign a subtask to a worker agent."""
171        if worker_id not in self.workers:
172            raise ValueError(f"Unknown worker: {worker_id}")
173 
174        self.workers[worker_id]["status"] = "busy"
175        self.workers[worker_id]["current_task"] = subtask.get("id")
176 
177        self._send(
178            AgentMessage(
179                sender=self.name,
180                receiver=worker_id,
181                message_type=MessageType.REQUEST,
182                content={"action": "execute_task", "task": subtask},
183                requires_response=True,
184                priority=subtask.get("priority", 0),
185            )
186        )
187 
188    def select_worker(self, subtask: Dict[str, Any]) -> str:
189        """Select the best available worker for a subtask.
190 
191        Use when: the supervisor needs capability-aware routing with
192        load-balancing (fewest completed tasks chosen first).
193        """
194        required_capability = subtask.get("type", "general")
195 
196        candidates = [
197            wid
198            for wid, info in self.workers.items()
199            if info["status"] == "available"
200            and required_capability in info["capabilities"]
201        ]
202 
203        if not candidates:
204            candidates = [
205                wid
206                for wid, info in self.workers.items()
207                if info["status"] == "available"
208            ]
209 
210        if not candidates:
211            raise ValueError("No available workers")
212 
213        return min(
214            candidates,
215            key=lambda w: self.workers[w]["metrics"]["tasks_completed"],
216        )
217 
218    def aggregate_results(
219        self, subtask_results: List[Dict[str, Any]]
220    ) -> Dict[str, Any]:
221        """Aggregate results from completed subtasks."""
222        summaries = [
223            r.get("summary", "")
224            for r in subtask_results
225            if r.get("success")
226        ]
227        successful = sum(
228            1 for r in subtask_results if r.get("success", False)
229        )
230        quality = successful / len(subtask_results) if subtask_results else 0.0
231 
232        return {
233            "results": subtask_results,
234            "summary": " | ".join(summaries),
235            "quality_score": quality,
236        }
237 
238    def run_workflow(self, task: Dict[str, Any]) -> Dict[str, Any]:
239        """Execute a complete workflow with supervision.
240 
241        Use when: running an end-to-end supervised pipeline that decomposes
242        a task, assigns subtasks, collects results, and aggregates them.
243 
244        Note: This is a synchronous simulation. Workers do not execute
245        asynchronously — each subtask is simulated inline. In production,
246        replace ``_simulate_worker_response`` with actual async worker
247        execution and message passing.
248        """
249        subtasks = self.decompose_task(task)
250 
251        results: List[Dict[str, Any]] = []
252        for subtask in subtasks:
253            worker = self.select_worker(subtask)
254            self.assign_task(subtask, worker)
255 
256            # Simulate worker executing and responding
257            response = self._simulate_worker_response(worker, subtask)
258            self.communication.send(
259                AgentMessage(
260                    sender=worker,
261                    receiver=self.name,
262                    message_type=MessageType.RESPONSE,
263                    content=response,
264                )
265            )
266            self.workers[worker]["status"] = "available"
267            self.workers[worker]["metrics"]["tasks_completed"] += 1
268 
269            messages = self.communication.receive(self.name)
270            for msg in messages:
271                if msg.message_type == MessageType.RESPONSE:
272                    results.append(msg.content)
273 
274        final_result = self.aggregate_results(results)
275 
276        return {
277            "task": task,
278            "subtask_results": results,
279            "final_result": final_result,
280            "success": final_result["quality_score"] >= 0.8,
281        }
282 
283    def _simulate_worker_response(
284        self, worker_id: str, subtask: Dict[str, Any]
285    ) -> Dict[str, Any]:
286        """Simulate a worker completing a subtask.
287 
288        In production, replace with actual agent execution that sends
289        the subtask to a worker process and awaits a real response.
290        """
291        return {
292            "success": True,
293            "summary": f"{worker_id} completed: {subtask.get('description', subtask.get('type', 'task'))}",
294            "worker": worker_id,
295            "subtask_type": subtask.get("type"),
296        }
297 
298    def _send(self, message: AgentMessage) -> None:
299        """Send message through the communication channel."""
300        self.communication.send(message)
301 
302 
303# ---------------------------------------------------------------------------
304# Handoff Protocol
305# ---------------------------------------------------------------------------
306 
307 
308class HandoffProtocol:
309    """Protocol for agent-to-agent handoffs.
310 
311    Use when: implementing peer-to-peer or swarm patterns where agents
312    transfer control and task state to one another.
313    """
314 
315    def __init__(self, communication: AgentCommunication) -> None:
316        self.communication = communication
317 
318    def create_handoff(
319        self,
320        from_agent: str,
321        to_agent: str,
322        context: Dict[str, Any],
323        reason: str,
324    ) -> AgentMessage:
325        """Create a handoff message with transferred context."""
326        return AgentMessage(
327            sender=from_agent,
328            receiver=to_agent,
329            message_type=MessageType.HANDOVER,
330            content={
331                "handoff_reason": reason,
332                "transferred_context": context,
333                "handoff_timestamp": time.time(),
334            },
335            priority=1,
336        )
337 
338    def accept_handoff(self, agent_id: str) -> Optional[AgentMessage]:
339        """Accept the first pending handoff for an agent, if any."""
340        messages = self.communication.receive(agent_id)
341 
342        for msg in messages:
343            if msg.message_type == MessageType.HANDOVER:
344                return msg
345 
346        return None
347 
348    def transfer_with_state(
349        self,
350        from_agent: str,
351        to_agent: str,
352        state: Dict[str, Any],
353        task: Dict[str, Any],
354    ) -> bool:
355        """Transfer task state from one agent to another.
356 
357        Use when: a handoff must carry full task state and progress so the
358        receiving agent can resume without re-deriving context.
359 
360        Returns True if the receiving agent acknowledged the handoff.
361        """
362        handoff = self.create_handoff(
363            from_agent=from_agent,
364            to_agent=to_agent,
365            context={
366                "task_state": state,
367                "task_details": task,
368                "progress": state.get("progress", 0),
369            },
370            reason="task_transfer",
371        )
372 
373        self.communication.send(handoff)
374 
375        # In production, replace sleep with async await + timeout
376        time.sleep(0.1)
377        ack = self.communication.receive(from_agent)
378 
379        return any(
380            m.message_type == MessageType.RESPONSE
381            and m.content.get("status") == "handoff_received"
382            for m in ack
383        )
384 
385 
386# ---------------------------------------------------------------------------
387# Consensus Mechanism
388# ---------------------------------------------------------------------------
389 
390 
391class ConsensusManager:
392    """Manager for multi-agent consensus building.
393 
394    Use when: multiple agents must vote on a decision and the system needs
395    weighted consensus that accounts for confidence and expertise rather
396    than naive majority voting.
397    """
398 
399    def __init__(self) -> None:
400        self.votes: Dict[str, List[Dict[str, Any]]] = {}
401        self.debates: Dict[str, List[Dict[str, Any]]] = {}
402 
403    def initiate_vote(
404        self, topic_id: str, agents: List[str], options: List[str]
405    ) -> None:
406        """Initiate a voting round on a topic."""
407        self.votes[topic_id] = [
408            {
409                "agent": agent,
410                "topic": topic_id,
411                "options": options,
412                "status": "pending",
413            }
414            for agent in agents
415        ]
416 
417    def submit_vote(
418        self,
419        topic_id: str,
420        agent_id: str,
421        selection: str,
422        confidence: float,
423    ) -> None:
424        """Submit a vote for a topic with a confidence weight."""
425        if topic_id not in self.votes:
426            raise ValueError(f"Unknown topic: {topic_id}")
427 
428        for vote in self.votes[topic_id]:
429            if vote["agent"] == agent_id:
430                vote["status"] = "cast"
431                vote["selection"] = selection
432                vote["confidence"] = confidence
433                break
434 
435    def calculate_weighted_consensus(self, topic_id: str) -> Dict[str, Any]:
436        """Calculate weighted consensus from cast votes.
437 
438        Use when: votes are in and the system needs to determine a winner
439        weighted by each agent's confidence rather than simple majority.
440        Weight = confidence * expertise_factor.
441        """
442        if topic_id not in self.votes:
443            raise ValueError(f"Unknown topic: {topic_id}")
444 
445        votes = [
446            v for v in self.votes[topic_id] if v.get("status") == "cast"
447        ]
448 
449        if not votes:
450            return {"status": "no_votes", "result": None}
451 
452        # Group by selection
453        selections: Dict[str, List[Dict[str, Any]]] = {}
454        for vote in votes:
455            selection = vote["selection"]
456            if selection not in selections:
457                selections[selection] = []
458            selections[selection].append(vote)
459 
460        # Calculate weighted score for each selection
461        results: Dict[str, Dict[str, Any]] = {}
462        for selection, selection_votes in selections.items():
463            weighted_sum = sum(v["confidence"] for v in selection_votes)
464            avg_confidence = (
465                weighted_sum / len(selection_votes) if selection_votes else 0.0
466            )
467            results[selection] = {
468                "weighted_score": weighted_sum,
469                "avg_confidence": avg_confidence,
470                "vote_count": len(selection_votes),
471            }
472 
473        winner = max(results.keys(), key=lambda s: results[s]["weighted_score"])
474 
475        return {
476            "status": "complete",
477            "result": winner,
478            "details": results,
479            "consensus_strength": (
480                results[winner]["weighted_score"] / len(votes) if votes else 0.0
481            ),
482        }
483 
484 
485# ---------------------------------------------------------------------------
486# Failure Handling
487# ---------------------------------------------------------------------------
488 
489 
490class AgentFailureHandler:
491    """Handler for agent failures in multi-agent systems.
492 
493    Use when: agents may fail and the system needs retry logic with
494    exponential backoff, circuit breakers, and automatic rerouting to
495    backup agents.
496    """
497 
498    def __init__(
499        self,
500        communication: AgentCommunication,
501        max_retries: int = 3,
502    ) -> None:
503        self.communication = communication
504        self.max_retries = max_retries
505        self.failure_counts: Dict[str, int] = {}
506        self.circuit_breakers: Dict[str, float] = {}  # agent -> unlock time
507 
508    def handle_failure(
509        self, agent_id: str, task_id: str, error: str
510    ) -> Dict[str, Any]:
511        """Handle a failure from an agent.
512 
513        Use when: an agent reports an error and the system must decide
514        whether to retry (with backoff) or reroute to a backup agent.
515        """
516        self.failure_counts[agent_id] = (
517            self.failure_counts.get(agent_id, 0) + 1
518        )
519 
520        if self.failure_counts[agent_id] >= self.max_retries:
521            self._activate_circuit_breaker(agent_id)
522            return {
523                "action": "reroute",
524                "reason": "circuit_breaker_activated",
525                "alternative": self._find_alternative_agent(agent_id),
526            }
527 
528        return {
529            "action": "retry",
530            "reason": error,
531            "retry_count": self.failure_counts[agent_id],
532            "delay": min(2 ** self.failure_counts[agent_id], 60),
533        }
534 
535    def _activate_circuit_breaker(self, agent_id: str) -> None:
536        """Temporarily disable an agent (1-minute cooldown)."""
537        self.circuit_breakers[agent_id] = time.time() + 60
538 
539    def _find_alternative_agent(self, failed_agent: str) -> str:
540        """Find an alternative agent to handle the task.
541 
542        In production, check agent capabilities and availability.
543        """
544        return "default_backup_agent"
545 
546    def is_available(self, agent_id: str) -> bool:
547        """Check if an agent is available (circuit breaker not active)."""
548        if agent_id in self.circuit_breakers:
549            if time.time() < self.circuit_breakers[agent_id]:
550                return False
551            del self.circuit_breakers[agent_id]
552            self.failure_counts[agent_id] = 0
553        return True
554 
555    def record_success(self, agent_id: str) -> None:
556        """Record a successful task completion, resetting failure count."""
557        self.failure_counts[agent_id] = 0
558 
559 
560# ---------------------------------------------------------------------------
561# Demo / CLI entry point
562# ---------------------------------------------------------------------------
563 
564 
565if __name__ == "__main__":
566    print("=== Multi-Agent Coordination Demo ===\n")
567 
568    # 1. Communication channel
569    comm = AgentCommunication()
570    print("1. Created communication channel")
571 
572    # 2. Supervisor pattern
573    supervisor = SupervisorAgent("supervisor", comm)
574    supervisor.register_worker("researcher", ["search", "analyze"])
575    supervisor.register_worker("writer", ["synthesize", "draft"])
576    print("2. Registered supervisor with 2 workers: researcher, writer")
577 
578    # 3. Handoff protocol
579    protocol = HandoffProtocol(comm)
580    handoff_msg = protocol.create_handoff(
581        from_agent="researcher",
582        to_agent="writer",
583        context={"findings": ["item1", "item2"]},
584        reason="research_complete",
585    )
586    comm.send(handoff_msg)
587    received = protocol.accept_handoff("writer")
588    print(
589        f"3. Handoff from researcher -> writer: "
590        f"{'accepted' if received else 'none pending'}"
591    )
592 
593    # 4. Consensus mechanism
594    consensus = ConsensusManager()
595    consensus.initiate_vote("best_approach", ["agent_a", "agent_b", "agent_c"], ["A", "B"])
596    consensus.submit_vote("best_approach", "agent_a", "A", confidence=0.9)
597    consensus.submit_vote("best_approach", "agent_b", "B", confidence=0.6)
598    consensus.submit_vote("best_approach", "agent_c", "A", confidence=0.8)
599    result = consensus.calculate_weighted_consensus("best_approach")
600    print(
601        f"4. Consensus result: {result['result']} "
602        f"(strength: {result['consensus_strength']:.2f})"
603    )
604 
605    # 5. Failure handling
606    handler = AgentFailureHandler(comm, max_retries=3)
607    action1 = handler.handle_failure("flaky_agent", "task_1", "timeout")
608    action2 = handler.handle_failure("flaky_agent", "task_1", "timeout")
609    action3 = handler.handle_failure("flaky_agent", "task_1", "timeout")
610    print(f"5. After 3 failures: action={action3['action']}")
611    print(f"   Agent available? {handler.is_available('flaky_agent')}")
612 
613    print("\n=== Demo Complete ===")
614