%pip install -qU anthropic

import asyncio
import itertools
from collections import Counter, defaultdict
 
import anthropic
 
MODEL = "claude-opus-4-8"  # swap for the newest Claude model available to you
client = anthropic.AsyncAnthropic()  # reads ANTHROPIC_API_KEY from env

class Hub:
    def __init__(self):
        self.inbox: dict[str, list[dict]] = defaultdict(list)
        self.event: dict[str, asyncio.Event] = defaultdict(asyncio.Event)
        self.status: dict[str, str] = {}
        self._ids = itertools.count(1)
 
    def register(self, name: str):
        self.status[name] = "active"
        _ = self.inbox[name], self.event[name]
 
    def new_name(self, prefix="helper") -> str:
        n = f"{prefix}{next(self._ids)}"
        self.register(n)
        return n
 
    def post(self, sender: str, recipients: list[str], content: str) -> list[str]:
        delivered = []
        for rid in recipients:
            if rid in self.status:
                self.inbox[rid].append({"from": sender, "content": content})
                self.event[rid].set()
                delivered.append(rid)
        return delivered
 
    def drain(self, name: str) -> list[dict]:
        msgs, self.inbox[name] = self.inbox[name], []
        self.event[name] = asyncio.Event()
        return msgs
 
    @staticmethod
    def render(msgs: list[dict]) -> str:
        if not msgs:
            return ""
        body = "\n".join(
            f'<agent-message from="{m["from"]}">\n{m["content"]}\n</agent-message>' for m in msgs
        )
        return f"\n\n[Messages received while you were working:]\n{body}"

SEND_MESSAGE = {
    "name": "send_message",
    "description": (
        "Send a message to one or more other agents. It will appear appended to their next tool "
        "result. This is the ONLY way to reach other agents — plain text in your turn goes nowhere."
    ),
    "input_schema": {
        "type": "object",
        "properties": {
            "recipient_ids": {"type": "array", "items": {"type": "string"}, "minItems": 1},
            "content": {"type": "string"},
        },
        "required": ["recipient_ids", "content"],
    },
}
WAIT_FOR_MESSAGE = {
    "name": "wait_for_message",
    "description": (
        "Block until another agent messages you. Note: messages also arrive automatically appended "
        "to the result of ANY other tool call, so only use this when you have nothing else to do."
    ),
    "input_schema": {"type": "object", "properties": {}},
}
BASE_TOOLS = [SEND_MESSAGE, WAIT_FOR_MESSAGE]

TRACE: dict[str, list[str]] = defaultdict(list)
 
 
def _snip(s, n=60):
    s = str(s).replace("\n", " ")
    return s if len(s) <= n else s[:n] + "…"
 
 
async def run_agent(
    hub: Hub,
    name: str,
    system: str,
    first_user_turn: str,
    tools: list = None,
    extra_dispatch=None,
    max_turns: int = 20,
) -> str:
    tools = tools or BASE_TOOLS
    extra_dispatch = extra_dispatch or {}
    messages = [{"role": "user", "content": first_user_turn}]
 
    try:
        for _ in range(max_turns):
            resp = await client.messages.create(
                model=MODEL,
                max_tokens=2048,
                system=system,
                tools=tools,
                messages=messages,
            )
            messages.append({"role": "assistant", "content": resp.content})
 
            if resp.stop_reason == "end_turn":
                hub.status[name] = "done"
                return "".join(getattr(b, "text", "") for b in resp.content)
            if resp.stop_reason != "tool_use":
                raise RuntimeError(f"unexpected stop_reason: {resp.stop_reason}")
 
            results = []
            for block in resp.content:
                if block.type != "tool_use":
                    continue
                TRACE[name].append(block.name)
                if block.name == "send_message":
                    rids = block.input["recipient_ids"]
                    delivered = hub.post(name, rids, block.input["content"])
                    unknown = [r for r in rids if r not in delivered]
                    out = f"delivered to {delivered}" + (f"; unknown: {unknown}" if unknown else "")
                elif block.name == "wait_for_message":
                    hub.status[name] = "idling"
                    try:
                        await asyncio.wait_for(hub.event[name].wait(), timeout=60)
                        out = "woke: new messages"
                    except TimeoutError:
                        out = "woke: 60s timeout"
                    hub.status[name] = "active"
                elif block.name in extra_dispatch:
                    out = await extra_dispatch[block.name](block)
                else:
                    out = f"error: no dispatch for {block.name}"
                print(f"  [{name}] {block.name}({_snip(block.input)}) → {_snip(out)}")
                results.append({"type": "tool_result", "tool_use_id": block.id, "content": out})
 
            inbox = hub.drain(name)
            for m in inbox:
                print(f"  [{name}] ← received from {m['from']}: {_snip(m['content'])}")
            if results:
                results[-1]["content"] += hub.render(inbox)  # ← the key line
            messages.append({"role": "user", "content": results})
 
        hub.status[name] = "done"
        return f"[{name} hit max_turns={max_turns}]"
    except Exception:
        hub.status[name] = "crashed"
        raise
 
 
def print_trace():
    for agent in sorted(TRACE):
        counts = Counter(TRACE[agent])
        print(f"  {agent}: " + ", ".join(f"{n}×{t}" for t, n in counts.most_common()))
    TRACE.clear()

TEAM_SYSTEM = "You are {name}, one of 3 agents working together (peers: {peers})."
 
TASKS = {
    "lead": "You are the lead. Introduce yourself to the others. Once everyone has introduced "
    "themselves, finish with a one-sentence summary of the team.",
    "helper1": "You are a backend engineer named Ada. Introduce yourself to the others, then wait "
    "for their replies.",
    "helper2": "You are a designer named Bo. Introduce yourself to the others, then wait for their "
    "replies.",
}
 
 
async def run_team() -> str:
    hub = Hub()
    names = list(TASKS)
    for n in names:
        hub.register(n)
 
    helper_tasks = [
        asyncio.create_task(
            run_agent(
                hub,
                n,
                system=TEAM_SYSTEM.format(name=n, peers=[p for p in names if p != n]),
                first_user_turn=TASKS[n],
            )
        )
        for n in names[1:]
    ]
    try:
        return await run_agent(
            hub,
            "lead",
            system=TEAM_SYSTEM.format(name="lead", peers=names[1:]),
            first_user_turn=TASKS["lead"],
        )
    finally:
        for t in helper_tasks:
            t.cancel()
        await asyncio.gather(*helper_tasks, return_exceptions=True)

answer = await run_team()
print(f"\n[lead final answer]\n{answer}\n\nTool-call summary:")
print_trace()

SLEEP = {
    "name": "sleep",
    "description": "Sleep for the given number of seconds, then return.",
    "input_schema": {
        "type": "object",
        "properties": {"seconds": {"type": "integer", "minimum": 0, "maximum": 10}},
        "required": ["seconds"],
    },
}
SUBAGENT_TOOLS = [
    {
        "name": "create_subagents",
        "description": "Spawn helper subagents. Returns immediately — helpers run concurrently in the "
        "background in parallel with you. Each gets base_instruction, optionally + "
        "per_subagent_instructions[i].",
        "input_schema": {
            "type": "object",
            "properties": {
                "base_instruction": {"type": "string"},
                "per_subagent_instructions": {
                    "type": "array",
                    "items": {"type": "string"},
                    "maxItems": 10,
                },
            },
            "required": ["base_instruction"],
        },
    },
    {
        "name": "get_status",
        "description": "Status of every helper (active / idling / done / crashed).",
        "input_schema": {"type": "object", "properties": {}},
    },
    {
        "name": "kill_subagents",
        "description": "Cancel running helpers you no longer need.",
        "input_schema": {
            "type": "object",
            "properties": {
                "subagent_ids": {"type": "array", "items": {"type": "string"}, "minItems": 1}
            },
            "required": ["subagent_ids"],
        },
    },
]
 
 
async def dispatch_sleep(block) -> str:
    s = int(block.input["seconds"])
    await asyncio.sleep(s)
    return f"slept {s}s"

async def run_spawn_lead() -> str:
    hub = Hub()
    hub.register("lead")
    helpers: dict[str, asyncio.Task] = {}
 
    async def _create(block):
        base = block.input["base_instruction"]
        per = block.input.get("per_subagent_instructions") or [""]
        spawned = []
        for suffix in per:
            h = hub.new_name()
            helpers[h] = asyncio.create_task(
                run_agent(
                    hub,
                    h,
                    system=f"You are {h}, a helper.",
                    first_user_turn=f"{base}\n\n{suffix}".strip(),
                    tools=[SLEEP, *BASE_TOOLS],
                    extra_dispatch={"sleep": dispatch_sleep},
                )
            )
            spawned.append(h)
        return f"spawned: {', '.join(spawned)}"
 
    async def _status(block):
        return "\n".join(f"{n}: {s}" for n, s in hub.status.items() if n != "lead") or "(none)"
 
    async def _kill(block):
        killed_ids, to_await = [], []
        for sid in block.input["subagent_ids"]:
            if sid in helpers and not helpers[sid].done():
                helpers[sid].cancel()
                hub.status[sid] = "done"
                to_await.append(helpers.pop(sid))
                killed_ids.append(sid)
        await asyncio.gather(*to_await, return_exceptions=True)
        return f"cancelled: {', '.join(killed_ids)}" if killed_ids else "no matching active helpers"
 
    try:
        return await run_agent(
            hub,
            "lead",
            system="You are the lead.",
            first_user_turn=(
                "Spawn three helper agents. Instruct each to sleep a different number of seconds "
                "(1, 2, 3), report back to you 'done, slept Ns', then wait for your further "
                "instructions. After that, check that they're running, collect all three reports, "
                "then dismiss all three helpers. Finish with a one-line summary."
            ),
            tools=[*SUBAGENT_TOOLS, *BASE_TOOLS],
            extra_dispatch={
                "create_subagents": _create,
                "get_status": _status,
                "kill_subagents": _kill,
            },
        )
    finally:
        for t in helpers.values():
            t.cancel()
        await asyncio.gather(*helpers.values(), return_exceptions=True)

answer = await run_spawn_lead()
print(f"\n[lead final answer]\n{answer}\n\nTool-call summary:")
print_trace()