【从零手写 ClaudeCode:learn-claude-code 项目实战笔记】(10)Team Protocols (团队协议)

0 阅读20分钟

第十章 Team Protocols (团队协议)

s01 > s02 > s03 > s04 > s05 > s06 | s07 > s08 > s09 > [ s10 ] s11 > s12

"本专栏基于开源项目 learn-claude-code 的官方文档。原文档非常硬核,为了方便像我一样的新手小白理解,我对文档进行了逐行精读,并加入了很多中文注释、大白话解释和踩坑记录。希望这套'咀嚼版'教程能帮你推开 AI Agent 开发的大门。"

项目地址:shareAI-lab/learn-claude-code: Bash is all you need - A nano Claude Code–like agent, built from 0 to 1

"队友之间要有统一的沟通规矩" -- 一个 request-response 模式驱动所有协商。

一、问题:s09 的消息太"随意"

s09 中队友之间能发消息,但全是"闲聊式"的,没有结构化流程。举两个会出问题的场景:

场景 1:关机 —— 直接杀线程会出事

在 s09 中,如果你关掉一个队友,只能这样粗暴处理:

# s09 只能这样
member["status"] = "shutdown"
# 或者直接等线程自然结束

问题在于:队友可能正在写文件、正在调 API,你直接标记 shutdown 会导致:

  • 文件写了一半,数据损坏
  • config.json 还显示它在 working,但线程已经结束了

正确做法应该是:先通知对方"我要关你了",等对方确认"我已经收尾完毕",再真正关闭。 这就是一个握手过程。

场景 2:计划审批 —— 高风险任务应该先过审

领导说"重构认证模块",s09 的队友接到 prompt 就直接开干。但重构是高风险操作,如果方案有问题,可能破坏整个系统。

正确做法应该是:先提交计划,等领导审查批准,再动手执行。

两个场景的共同模式

你仔细看这两个场景,结构完全一样:

场景 1(关机):  领导发请求 → 队友审批 → 用某种方式关联请求和响应
场景 2(计划):  队友发请求 → 领导审批 → 用某种方式关联请求和响应

区别只是谁发起、谁审批不同,但底层模式完全一致:一方发请求,另一方用 approve/reject 响应,两边通过同一个 ID 关联起来。

这就是 s10 要做的:把这种 request-response 模式抽象成一套通用协议,同时覆盖关机和计划审批两个场景。

二、解决方案

先看全局架构图:

Shutdown Protocol(关机协议)        Plan Approval Protocol(计划审批协议)
===========================        ======================================

Lead              Teammate         Teammate              Lead
  |                  |               |                    |
  |--shutdown_req--->|               |--plan_req--------->|
  | {req_id:"abc"}   |               | {req_id:"xyz"}     |
  |                  |               |                    |
  |<--shutdown_resp--|               |<--plan_resp--------|
  | {req_id:"abc",   |               | {req_id:"xyz",     |
  |  approve:true}   |               |  approve:true}     |

共享的状态机(FSM):
  [pending] --approve--> [approved]
  [pending] --reject---> [rejected]

请求追踪器(内存字典):
  shutdown_requests = {req_id: {target, status}}
  plan_requests     = {req_id: {from, plan, status}}

"协议"到底是什么?

在 s09 中,send_message 就是随便发一句话:

{"type": "message", "from": "alice", "content": "我干完了"}

在 s10 中,协议消息多了三样东西

{
  "type": "shutdown_response",
  "from": "alice",
  "content": "All files saved.",
  "request_id": "abc12345",1. 关联 ID
  "approve": true2. 审批结果
}

加上代码层面的支撑:

    1. 内存字典追踪状态:shutdown_requests["abc12345"]["status"] = "approved"
    1. 状态机保证只能走 pending → approved | rejected 的流程

所以"协议" = 固定的 JSON 格式 + 固定的交互流程 + 状态追踪。 三者缺一不可。

核心三要素

  1. request_id — 每个请求一个唯一 ID,用来把请求和响应关联起来
  2. FSM(有限状态机) — 统一的状态流转:pending → approved | rejected
  3. Tracker — 内存字典,追踪每个请求的当前状态

三、工作原理

3.1 为什么要用 request_id?

在 s09 中,消息之间没有任何关联。领导给 alice 发消息"请关闭",alice 回复"好的",但这两条消息在代码层面是独立的,没有任何字段把它们关联起来。

问题来了: 如果同时有多个请求呢?

lead → alice: "请关闭"(请求 A)
lead → bob: "请关闭"(请求 B)
alice → lead: "好的"    ← 这个"好的"是回复请求 A 还是请求 B?
bob → lead: "不行"      ← 这个"不行"是回复请求 A 还是请求 B?

request_id 就是解决这个问题的。 每个请求带一个唯一 ID,响应必须引用同一个 ID:

lead → alice: "请关闭" {request_id: "abc12345"}(请求 A)
lead → bob: "请关闭" {request_id: "def67890"}(请求 B)
alice → lead: "好的" {request_id: "abc12345", approve: true}   ← 明确是回复请求 A
bob → lead: "不行" {request_id: "def67890", approve: false}     ← 明确是回复请求 B

代码中的追踪器:

shutdown_requests = {}   # 追踪关机请求
plan_requests = {}       # 追踪计划审批请求
_tracker_lock = threading.Lock()  # 线程安全锁(多线程同时读写需要加锁)

每个请求以 request_id 为 key,存入内存字典:

# 关机请求被创建时
shutdown_requests = {
    "abc12345": {"target": "alice", "status": "pending"},
}

# 关机请求被响应后
shutdown_requests = {
    "abc12345": {"target": "alice", "status": "approved"},  # 状态更新
}

# 计划审批请求
plan_requests = {
    "xyz54321": {"from": "alice", "plan": "重构认证模块,分三步...", "status": "pending"},
}

为什么用内存字典而不是文件? 因为请求状态是临时的,请求完成后就不需要了。用内存更快,进程退出时自动清理。而 config.json(团队名册)需要持久化,所以用文件。

3.2 _tracker_lock 是什么?为什么需要它?

_tracker_lock = threading.Lock()

因为队友跑在独立线程里,可能同时修改 shutdown_requestsplan_requests。Python 的字典在多线程同时写入时可能出错,所以用 Lock(锁)保证同一时刻只有一个线程在操作字典:

with _tracker_lock:
    shutdown_requests[req_id] = {"target": teammate, "status": "pending"}

with _tracker_lock 的意思是:进入这段代码前先"拿锁",执行完后"释放锁"。如果另一个线程同时想拿锁,它会等待,直到锁被释放。

3.3 关机协议完整流程

关机协议涉及两个代码执行位置

位置 1:领导的 agent_loop(主线程)
        → 调用 TOOL_HANDLERS 里的 lambda
        → 最终调用 handle_shutdown_request()

位置 2:队友的 _teammate_loop(子线程)
        → 调用 self._exec()
        → 处理 shutdown_response 工具

下面逐步走完整个流程。

第一步:领导发起关机请求

领导的 LLM 决定要关掉 alice,调用 shutdown_request 工具:

# 领导的工具分发(TOOL_HANDLERS)
"shutdown_request": lambda **kw: handle_shutdown_request(kw["teammate"])

实际执行的函数:

def handle_shutdown_request(teammate: str) -> str:
    # 1. 生成唯一 request_id(取 UUID 的前 8 位,够用且短)
    req_id = str(uuid.uuid4())[:8]

    # 2. 记录到追踪器,状态为 pending
    with _tracker_lock:
        shutdown_requests[req_id] = {"target": teammate, "status": "pending"}

    # 3. 发送 shutdown_request 消息到队友的收件箱
    #    注意 extra 里带上了 request_id
    BUS.send(
        "lead",              # 发送者
        teammate,            # 收件人,比如 "alice"
        "Please shut down gracefully.",  # 消息内容
        "shutdown_request",  # 消息类型
        {"request_id": req_id},  # extra:把 request_id 附带进去
    )
    return f"Shutdown request {req_id} sent to '{teammate}' (status: pending)"

send() 方法里 msg.update(extra) 会把 request_id 合并进消息,最终写入 alice.jsonl 的消息:

{
  "type": "shutdown_request",
  "from": "lead",
  "content": "Please shut down gracefully.",
  "timestamp": 1710912345.67,
  "request_id": "abc12345"
}

同时内存中的追踪器状态:

shutdown_requests = {
    "abc12345": {"target": "alice", "status": "pending"}
}
第二步:队友读到请求,决定批准还是拒绝

队友的 _teammate_loop 在每轮循环开头会检查收件箱:

for _ in range(50):
    inbox = BUS.read_inbox(name)  # 读取 alice 的收件箱
    for msg in inbox:
        # 收到的消息注入到 messages 数组,LLM 会看到
        messages.append({"role": "user", "content": json.dumps(msg)})

    response = client.messages.create(...)  # LLM 看到消息后做决策
    # ...

LLM 看到收件箱里的 shutdown_request 消息后,会判断自己的工作状态。如果它觉得自己可以关闭了,就调用 shutdown_response 工具。

这个工具在 self._exec() 里处理:

def _exec(self, sender: str, tool_name: str, args: dict) -> str:
    # ...
    if tool_name == "shutdown_response":
        req_id = args["request_id"]       # 拿到 request_id
        approve = args["approve"]         # 拿到 approve(true 或 false)

        # 1. 更新追踪器状态
        with _tracker_lock:
            if req_id in shutdown_requests:
                shutdown_requests[req_id]["status"] = "approved" if approve else "rejected"

        # 2. 发送响应消息回 lead
        BUS.send(
            sender,  # "alice"
            "lead",
            args.get("reason", ""),  # LLM 可以填原因,比如 "All files saved"
            "shutdown_response",
            {"request_id": req_id, "approve": approve},
        )
        return f"Shutdown {'approved' if approve else 'rejected'}"

写入 lead.jsonl 的消息:

{
  "type": "shutdown_response",
  "from": "alice",
  "content": "All files saved, ready to shutdown.",
  "timestamp": 1710912346.67,
  "request_id": "abc12345",
  "approve": true
}

同时追踪器更新:

shutdown_requests = {
    "abc12345": {"target": "alice", "status": "approved"}
}
第三步:队友优雅退出

_teammate_loop 中,有一个 should_exit 标志位控制退出:

def _teammate_loop(self, name, role, prompt):
    should_exit = False
    for _ in range(50):
        inbox = BUS.read_inbox(name)
        for msg in inbox:
            messages.append({"role": "user", "content": json.dumps(msg)})

        if should_exit:
            break  # 上一轮标记了退出,这一轮直接退出循环

        response = client.messages.create(...)
        # ...

        results = []
        for block in response.content:
            if block.type == "tool_use":
                output = self._exec(name, block.name, block.input)
                print(f"  [{name}] {block.name}: {str(output)[:120]}")
                results.append(...)

                # 关键:检测到队友调用了 shutdown_response 且 approve=true
                if block.name == "shutdown_response" and block.input.get("approve"):
                    should_exit = True  # 标记退出,但不立刻 break
                    # 为什么要延迟?因为当前轮次可能还有其他工具调用需要处理完

        messages.append({"role": "user", "content": results})

    # 循环结束后,根据退出原因设置最终状态
    member = self._find_member(name)
    if member:
        member["status"] = "shutdown" if should_exit else "idle"
        self._save_config()

should_exit 的延迟退出设计: LLM 一轮可能同时调用多个工具,比如先 write_file 保存文件,再 shutdown_response 批准关闭。如果收到 shutdown_response 就立刻 break,write_file 的结果就丢了。所以标记退出后让当前轮次执行完,下一轮再退出。

总结:关机的完整调用链
用户输入: "请关闭 alice"

Lead agent loop:
  LLM 决定调用 shutdown_request(teammate="alice")
    → TOOL_HANDLERS 匹配到 lambda
    → 调用 handle_shutdown_request("alice")
    → 生成 request_id "abc12345"
    → shutdown_requests["abc12345"] = {"target": "alice", "status": "pending"}
    → BUS.send() 写入 alice.jsonl
    → 返回 "Shutdown request abc12345 sent to 'alice' (status: pending)"

Alice _teammate_loop(独立线程):
  read_inbox("alice") 读到 shutdown_request
  → 注入到 messages,LLM 看到消息
  → LLM 判断"我还没写完",决定拒绝
  → 调用 shutdown_response(request_id="abc12345", approve=false, reason="还在写文件")
    → self._exec() 匹配到 shutdown_response 分支
    → shutdown_requests["abc12345"]["status"] = "rejected"
    → BUS.send() 写入 lead.jsonl
  → should_exit 没被设为 True,继续工作

Lead agent loop:
  read_inbox("lead") 读到 shutdown_response (approve=false)
  → LLM 看到"alice 拒绝了"
  → LLM 决定等一会儿再发

  (等了一会儿,Lead 再次发送)
  → handle_shutdown_request("alice"),新的 request_id "def67890"

Alice _teammate_loop:
  read_inbox("alice") 读到新的 shutdown_request
  → LLM 判断"我写完了"
  → 调用 shutdown_response(request_id="def67890", approve=true)
  → should_exit = True
  → 当前轮次执行完后退出循环
  → status = "shutdown"

3.4 计划审批协议(Plan Approval Protocol)

和关机协议完全相同的模式,只是方向相反:队友发起请求,领导审批。

第一步:队友提交计划

当队友要执行高风险任务(比如重构代码)时,先调用 plan_approval 工具提交计划:

# self._exec() 中
if tool_name == "plan_approval":
    plan_text = args.get("plan", "")
    req_id = str(uuid.uuid4())[:8]  # 队友自己生成 request_id

    # 1. 记录到计划追踪器
    with _tracker_lock:
        plan_requests[req_id] = {"from": sender, "plan": plan_text, "status": "pending"}

    # 2. 发送计划给 lead
    BUS.send(
        sender,     # "alice"
        "lead",
        plan_text,  # 计划内容
        "plan_approval_response",
        {"request_id": req_id, "plan": plan_text},
    )
    return f"Plan submitted (request_id={req_id}). Waiting for lead approval."

写入 lead.jsonl 的消息:

{
  "type": "plan_approval_response",
  "from": "alice",
  "content": "重构认证模块,分三步:1. 提取接口 2. 实现新方案 3. 迁移旧调用",
  "timestamp": 1710912345.67,
  "request_id": "xyz54321",
  "plan": "重构认证模块,分三步:1. 提取接口 2. 实现新方案 3. 迁移旧调用"
}

此时队友应该等待领导的审批结果,而不是直接开始干。 怎么等待?LLM 一轮结束后,下一轮循环开头会 read_inbox,看看有没有领导的回复。

第二步:领导审批

领导的 LLM 看到 alice 提交的计划后,调用 plan_approval 工具进行审批:

# 领导的 TOOL_HANDLERS 中
"plan_approval": lambda **kw: handle_plan_review(
    kw["request_id"], kw["approve"], kw.get("feedback", "")
)

实际执行的函数:

def handle_plan_review(request_id: str, approve: bool, feedback: str = "") -> str:
    # 1. 通过 request_id 找到对应的计划请求
    with _tracker_lock:
        req = plan_requests.get(request_id)
    if not req:
        return f"Error: Unknown plan request_id '{request_id}'"

    # 2. 更新状态
    with _tracker_lock:
        req["status"] = "approved" if approve else "rejected"

    # 3. 发送审批结果回队友
    BUS.send(
        "lead",
        req["from"],     # "alice"(从 plan_requests 里取出来的)
        feedback,        # 领导的反馈意见
        "plan_approval_response",
        {"request_id": request_id, "approve": approve, "feedback": feedback},
    )
    return f"Plan {req['status']} for '{req['from']}'"

写入 alice.jsonl 的消息:

{
  "type": "plan_approval_response",
  "from": "lead",
  "content": "步骤 2 风险太高,先做原型验证",
  "timestamp": 1710912346.67,
  "request_id": "xyz54321",
  "approve": false,
  "feedback": "步骤 2 风险太高,先做原型验证"
}
第三步:队友收到审批结果

alice 的 _teammate_loop 下一轮 read_inbox 读到这条消息,注入到 messages。LLM 看到 approve: false 和反馈意见,决定调整计划再提交,或者接受拒绝停止工作。

3.5 两个协议的对比

                     关机协议                      计划审批协议

谁发起请求           Lead                         Teammate
谁审批               Teammate                     Lead
request_id 谁生成    Lead(handle_shutdown_request) Teammate(plan_approval 工具)
追踪器               shutdown_requests              plan_requests
发起方工具           shutdown_request(领导)        plan_approval(队友)
审批方工具           shutdown_response(队友)       plan_approval(领导)
FSM                  pending → approved/rejected     pending → approved/rejected

代码体现:

# 关机:领导发起,记录到 shutdown_requests
shutdown_requests[req_id] = {"target": teammate, "status": "pending"}

# 计划:队友发起,记录到 plan_requests
plan_requests[req_id] = {"from": sender, "plan": plan_text, "status": "pending"}

# 两者审批时的状态更新完全一样
status = "approved" if approve else "rejected"

3.6 领导和队友的工具对比

s10 相对 s09,两边各新增了工具:

领导的工具(12 个,比 s09 多 3 个):

TOOL_HANDLERS = {
    # s09 已有的工具
    "bash": ...,
    "read_file": ...,
    "write_file": ...,
    "edit_file": ...,
    "spawn_teammate": ...,
    "list_teammates": ...,
    "send_message": ...,
    "read_inbox": ...,
    "broadcast": ...,
    # s10 新增
    "shutdown_request":  lambda **kw: handle_shutdown_request(kw["teammate"]),
    "shutdown_response": lambda **kw: _check_shutdown_status(kw.get("request_id", "")),
    "plan_approval":     lambda **kw: handle_plan_review(
                           kw["request_id"], kw["approve"], kw.get("feedback", "")),
}

注意:领导的 shutdown_response 工具 不是用来审批的,而是用来查询关机请求的状态(通过 _check_shutdown_status)。审批是由队友调用 shutdown_response 完成的。

队友的工具(8 个,比 s09 多 2 个):

def _teammate_tools(self) -> list:
    return [
        # s09 已有的工具
        {"name": "bash", ...},
        {"name": "read_file", ...},
        {"name": "write_file", ...},
        {"name": "edit_file", ...},
        {"name": "send_message", ...},
        {"name": "read_inbox", ...},
        # s10 新增
        {"name": "shutdown_response", ...},  # 审批关机请求
        {"name": "plan_approval", ...},       # 提交计划等待审批
    ]

3.7 LLM 为什么知道按协议格式回复?

因为工具的 input_schema + system prompt 共同约束了 LLM 的行为。

第一步:工具 schema 规定了"必须填什么参数"

# shutdown_response 工具定义
{
    "name": "shutdown_response",
    "description": "Respond to a shutdown request. Approve to shut down, reject to keep working.",
    "input_schema": {
        "type": "object",
        "properties": {
            "request_id": {"type": "string"},      # 必须填 request_id
            "approve": {"type": "boolean"},         # 必须填 true/false
            "reason": {"type": "string"}            # 可选填原因
        },
        "required": ["request_id", "approve"]
    }
}

LLM 调用这个工具时,必须提供 request_idapprove,否则 API 会返回错误。

第二步:system prompt 告诉 LLM"什么时候该用什么工具"

sys_prompt = (
    f"You are '{name}', role: {role}, at {WORKDIR}. "
    f"Submit plans via plan_approval before major work. "
    f"Respond to shutdown_request with shutdown_response."
)

"Respond to shutdown_request with shutdown_response" 就是告诉 LLM:

"当你在收件箱里看到一条 type: shutdown_request 的消息时,用 shutdown_response 工具回复,不要用普通的 send_message。"

如果没有这句话,LLM 可能会用普通消息回复:

# ❌ 错误做法:用 send_message 回复
send_message(to="lead", content="好的我关了")
# 这样 lead 收到的只是一条普通 message,没有 request_id,没有 approve 字段

有了这句话,LLM 就知道用协议工具回复:

# ✅ 正确做法:用 shutdown_response 回复
shutdown_response(request_id="abc12345", approve=true, reason="done")
# 这样 lead 收到的是结构化响应,可以精确关联和处理

四、相对 s09 的变更

组件之前 (s09)之后 (s10)
领导工具912 (+3)
队友工具68 (+2)
关机请求-响应握手
计划门控提交/审查与审批
关联每个请求一个 request_id
FSMpending → approved/rejected
新增全局变量-shutdown_requests、plan_requests、_tracker_lock

五、关键代码对比

s09 关机 vs s10 关机

s09:直接改状态,没有通知、没有确认

# 一行代码搞定,但也埋下了隐患
member["status"] = "shutdown"

s10:完整的握手流程

# 1. 生成 request_id
req_id = str(uuid.uuid4())[:8]

# 2. 记录状态
with _tracker_lock:
    shutdown_requests[req_id] = {"target": teammate, "status": "pending"}

# 3. 发送请求(队友的 agent loop 在另一线程中处理)
BUS.send("lead", teammate, "Please shut down.", "shutdown_request", {"request_id": req_id})

# 4. 等待队友响应(read_inbox 读到 response)
# 5. 队友更新状态
shutdown_requests[req_id]["status"] = "approved"

# 6. 队友标记 should_exit = True,优雅退出

一句话总结:s09 是"随便发消息",s10 是"有问必答、有 ID 可追踪的结构化对话"。

六、试一试

cd learn-claude-code
python agents/s10_team_protocols.py

试试这些 prompt(英文 prompt 对 LLM 效果更好,也可以用中文):

  1. Spawn alice as a coder. Then request her shutdown.
  2. List teammates to see alice's status after shutdown approval
  3. Spawn bob with a risky refactoring task. Review and reject his plan.
  4. Spawn charlie, have him submit a plan, then approve it.
  5. 输入 /team 监控状态

七、完整代码

#!/usr/bin/env python3
import json
import os
import subprocess
import threading
import time
import uuid
from pathlib import Path

from anthropic import Anthropic
from dotenv import load_dotenv

load_dotenv(override=True)
if os.getenv("ANTHROPIC_BASE_URL"):
    os.environ.pop("ANTHROPIC_AUTH_TOKEN", None)

WORKDIR = Path.cwd()
client = Anthropic(base_url=os.getenv("ANTHROPIC_BASE_URL"))
MODEL = os.environ["MODEL_ID"]
TEAM_DIR = WORKDIR / ".team"
INBOX_DIR = TEAM_DIR / "inbox"

SYSTEM = f"You are a team lead at {WORKDIR}. Manage teammates with shutdown and plan approval protocols."

VALID_MSG_TYPES = {
    "message",
    "broadcast",
    "shutdown_request",
    "shutdown_response",
    "plan_approval_response",
}

# -- Request trackers: correlate by request_id --
shutdown_requests = {}
plan_requests = {}
_tracker_lock = threading.Lock()


# -- MessageBus: JSONL inbox per teammate --
class MessageBus:
    def __init__(self, inbox_dir: Path):
        self.dir = inbox_dir
        self.dir.mkdir(parents=True, exist_ok=True)

    def send(self, sender: str, to: str, content: str,
             msg_type: str = "message", extra: dict = None) -> str:
        if msg_type not in VALID_MSG_TYPES:
            return f"Error: Invalid type '{msg_type}'. Valid: {VALID_MSG_TYPES}"
        msg = {
            "type": msg_type,
            "from": sender,
            "content": content,
            "timestamp": time.time(),
        }
        if extra:
            msg.update(extra)
        inbox_path = self.dir / f"{to}.jsonl"
        with open(inbox_path, "a") as f:
            f.write(json.dumps(msg) + "\n")
        return f"Sent {msg_type} to {to}"

    def read_inbox(self, name: str) -> list:
        inbox_path = self.dir / f"{name}.jsonl"
        if not inbox_path.exists():
            return []
        messages = []
        for line in inbox_path.read_text().strip().splitlines():
            if line:
                messages.append(json.loads(line))
        inbox_path.write_text("")
        return messages

    def broadcast(self, sender: str, content: str, teammates: list) -> str:
        count = 0
        for name in teammates:
            if name != sender:
                self.send(sender, name, content, "broadcast")
                count += 1
        return f"Broadcast to {count} teammates"


BUS = MessageBus(INBOX_DIR)


# -- TeammateManager with shutdown + plan approval --
class TeammateManager:
    def __init__(self, team_dir: Path):
        self.dir = team_dir
        self.dir.mkdir(exist_ok=True)
        self.config_path = self.dir / "config.json"
        self.config = self._load_config()
        self.threads = {}

    def _load_config(self) -> dict:
        if self.config_path.exists():
            return json.loads(self.config_path.read_text())
        return {"team_name": "default", "members": []}

    def _save_config(self):
        self.config_path.write_text(json.dumps(self.config, indent=2))

    def _find_member(self, name: str) -> dict:
        for m in self.config["members"]:
            if m["name"] == name:
                return m
        return None

    def spawn(self, name: str, role: str, prompt: str) -> str:
        member = self._find_member(name)
        if member:
            if member["status"] not in ("idle", "shutdown"):
                return f"Error: '{name}' is currently {member['status']}"
            member["status"] = "working"
            member["role"] = role
        else:
            member = {"name": name, "role": role, "status": "working"}
            self.config["members"].append(member)
        self._save_config()
        thread = threading.Thread(
            target=self._teammate_loop,
            args=(name, role, prompt),
            daemon=True,
        )
        self.threads[name] = thread
        thread.start()
        return f"Spawned '{name}' (role: {role})"

    def _teammate_loop(self, name: str, role: str, prompt: str):
        sys_prompt = (
            f"You are '{name}', role: {role}, at {WORKDIR}. "
            f"Submit plans via plan_approval before major work. "
            f"Respond to shutdown_request with shutdown_response."
        )
        messages = [{"role": "user", "content": prompt}]
        tools = self._teammate_tools()
        should_exit = False
        for _ in range(50):
            inbox = BUS.read_inbox(name)
            for msg in inbox:
                messages.append({"role": "user", "content": json.dumps(msg)})
            if should_exit:
                break
            try:
                response = client.messages.create(
                    model=MODEL,
                    system=sys_prompt,
                    messages=messages,
                    tools=tools,
                    max_tokens=8000,
                )
            except Exception:
                break
            messages.append({"role": "assistant", "content": response.content})
            if response.stop_reason != "tool_use":
                break
            results = []
            for block in response.content:
                if block.type == "tool_use":
                    output = self._exec(name, block.name, block.input)
                    print(f"  [{name}] {block.name}: {str(output)[:120]}")
                    results.append({
                        "type": "tool_result",
                        "tool_use_id": block.id,
                        "content": str(output),
                    })
                    if block.name == "shutdown_response" and block.input.get("approve"):
                        should_exit = True
            messages.append({"role": "user", "content": results})
        member = self._find_member(name)
        if member:
            member["status"] = "shutdown" if should_exit else "idle"
            self._save_config()

    def _exec(self, sender: str, tool_name: str, args: dict) -> str:
        # these base tools are unchanged from s02
        if tool_name == "bash":
            return _run_bash(args["command"])
        if tool_name == "read_file":
            return _run_read(args["path"])
        if tool_name == "write_file":
            return _run_write(args["path"], args["content"])
        if tool_name == "edit_file":
            return _run_edit(args["path"], args["old_text"], args["new_text"])
        if tool_name == "send_message":
            return BUS.send(sender, args["to"], args["content"], args.get("msg_type", "message"))
        if tool_name == "read_inbox":
            return json.dumps(BUS.read_inbox(sender), indent=2)
        if tool_name == "shutdown_response":
            req_id = args["request_id"]
            approve = args["approve"]
            with _tracker_lock:
                if req_id in shutdown_requests:
                    shutdown_requests[req_id]["status"] = "approved" if approve else "rejected"
            BUS.send(
                sender, "lead", args.get("reason", ""),
                "shutdown_response", {"request_id": req_id, "approve": approve},
            )
            return f"Shutdown {'approved' if approve else 'rejected'}"
        if tool_name == "plan_approval":
            plan_text = args.get("plan", "")
            req_id = str(uuid.uuid4())[:8]
            with _tracker_lock:
                plan_requests[req_id] = {"from": sender, "plan": plan_text, "status": "pending"}
            BUS.send(
                sender, "lead", plan_text, "plan_approval_response",
                {"request_id": req_id, "plan": plan_text},
            )
            return f"Plan submitted (request_id={req_id}). Waiting for lead approval."
        return f"Unknown tool: {tool_name}"

    def _teammate_tools(self) -> list:
        # these base tools are unchanged from s02
        return [
            {"name": "bash", "description": "Run a shell command.",
             "input_schema": {"type": "object", "properties": {"command": {"type": "string"}}, "required": ["command"]}},
            {"name": "read_file", "description": "Read file contents.",
             "input_schema": {"type": "object", "properties": {"path": {"type": "string"}}, "required": ["path"]}},
            {"name": "write_file", "description": "Write content to file.",
             "input_schema": {"type": "object", "properties": {"path": {"type": "string"}, "content": {"type": "string"}}, "required": ["path", "content"]}},
            {"name": "edit_file", "description": "Replace exact text in file.",
             "input_schema": {"type": "object", "properties": {"path": {"type": "string"}, "old_text": {"type": "string"}, "new_text": {"type": "string"}}, "required": ["path", "old_text", "new_text"]}},
            {"name": "send_message", "description": "Send message to a teammate.",
             "input_schema": {"type": "object", "properties": {"to": {"type": "string"}, "content": {"type": "string"}, "msg_type": {"type": "string", "enum": list(VALID_MSG_TYPES)}}, "required": ["to", "content"]}},
            {"name": "read_inbox", "description": "Read and drain your inbox.",
             "input_schema": {"type": "object", "properties": {}}},
            {"name": "shutdown_response", "description": "Respond to a shutdown request. Approve to shut down, reject to keep working.",
             "input_schema": {"type": "object", "properties": {"request_id": {"type": "string"}, "approve": {"type": "boolean"}, "reason": {"type": "string"}}, "required": ["request_id", "approve"]}},
            {"name": "plan_approval", "description": "Submit a plan for lead approval. Provide plan text.",
             "input_schema": {"type": "object", "properties": {"plan": {"type": "string"}}, "required": ["plan"]}},
        ]

    def list_all(self) -> str:
        if not self.config["members"]:
            return "No teammates."
        lines = [f"Team: {self.config['team_name']}"]
        for m in self.config["members"]:
            lines.append(f"  {m['name']} ({m['role']}): {m['status']}")
        return "\n".join(lines)

    def member_names(self) -> list:
        return [m["name"] for m in self.config["members"]]


TEAM = TeammateManager(TEAM_DIR)


# -- Base tool implementations (these base tools are unchanged from s02) --
def _safe_path(p: str) -> Path:
    path = (WORKDIR / p).resolve()
    if not path.is_relative_to(WORKDIR):
        raise ValueError(f"Path escapes workspace: {p}")
    return path


def _run_bash(command: str) -> str:
    dangerous = ["rm -rf /", "sudo", "shutdown", "reboot"]
    if any(d in command for d in dangerous):
        return "Error: Dangerous command blocked"
    try:
        r = subprocess.run(
            command, shell=True, cwd=WORKDIR,
            capture_output=True, text=True, timeout=120,
        )
        out = (r.stdout + r.stderr).strip()
        return out[:50000] if out else "(no output)"
    except subprocess.TimeoutExpired:
        return "Error: Timeout (120s)"


def _run_read(path: str, limit: int = None) -> str:
    try:
        lines = _safe_path(path).read_text().splitlines()
        if limit and limit < len(lines):
            lines = lines[:limit] + [f"... ({len(lines) - limit} more)"]
        return "\n".join(lines)[:50000]
    except Exception as e:
        return f"Error: {e}"


def _run_write(path: str, content: str) -> str:
    try:
        fp = _safe_path(path)
        fp.parent.mkdir(parents=True, exist_ok=True)
        fp.write_text(content)
        return f"Wrote {len(content)} bytes"
    except Exception as e:
        return f"Error: {e}"


def _run_edit(path: str, old_text: str, new_text: str) -> str:
    try:
        fp = _safe_path(path)
        c = fp.read_text()
        if old_text not in c:
            return f"Error: Text not found in {path}"
        fp.write_text(c.replace(old_text, new_text, 1))
        return f"Edited {path}"
    except Exception as e:
        return f"Error: {e}"


# -- Lead-specific protocol handlers --
def handle_shutdown_request(teammate: str) -> str:
    req_id = str(uuid.uuid4())[:8]
    with _tracker_lock:
        shutdown_requests[req_id] = {"target": teammate, "status": "pending"}
    BUS.send(
        "lead", teammate, "Please shut down gracefully.",
        "shutdown_request", {"request_id": req_id},
    )
    return f"Shutdown request {req_id} sent to '{teammate}' (status: pending)"


def handle_plan_review(request_id: str, approve: bool, feedback: str = "") -> str:
    with _tracker_lock:
        req = plan_requests.get(request_id)
    if not req:
        return f"Error: Unknown plan request_id '{request_id}'"
    with _tracker_lock:
        req["status"] = "approved" if approve else "rejected"
    BUS.send(
        "lead", req["from"], feedback, "plan_approval_response",
        {"request_id": request_id, "approve": approve, "feedback": feedback},
    )
    return f"Plan {req['status']} for '{req['from']}'"


def _check_shutdown_status(request_id: str) -> str:
    with _tracker_lock:
        return json.dumps(shutdown_requests.get(request_id, {"error": "not found"}))


# -- Lead tool dispatch (12 tools) --
TOOL_HANDLERS = {
    "bash":              lambda **kw: _run_bash(kw["command"]),
    "read_file":         lambda **kw: _run_read(kw["path"], kw.get("limit")),
    "write_file":        lambda **kw: _run_write(kw["path"], kw["content"]),
    "edit_file":         lambda **kw: _run_edit(kw["path"], kw["old_text"], kw["new_text"]),
    "spawn_teammate":    lambda **kw: TEAM.spawn(kw["name"], kw["role"], kw["prompt"]),
    "list_teammates":    lambda **kw: TEAM.list_all(),
    "send_message":      lambda **kw: BUS.send("lead", kw["to"], kw["content"], kw.get("msg_type", "message")),
    "read_inbox":        lambda **kw: json.dumps(BUS.read_inbox("lead"), indent=2),
    "broadcast":         lambda **kw: BUS.broadcast("lead", kw["content"], TEAM.member_names()),
    "shutdown_request":  lambda **kw: handle_shutdown_request(kw["teammate"]),
    "shutdown_response": lambda **kw: _check_shutdown_status(kw.get("request_id", "")),
    "plan_approval":     lambda **kw: handle_plan_review(kw["request_id"], kw["approve"], kw.get("feedback", "")),
}

# these base tools are unchanged from s02
TOOLS = [
    {"name": "bash", "description": "Run a shell command.",
     "input_schema": {"type": "object", "properties": {"command": {"type": "string"}}, "required": ["command"]}},
    {"name": "read_file", "description": "Read file contents.",
     "input_schema": {"type": "object", "properties": {"path": {"type": "string"}, "limit": {"type": "integer"}}, "required": ["path"]}},
    {"name": "write_file", "description": "Write content to file.",
     "input_schema": {"type": "object", "properties": {"path": {"type": "string"}, "content": {"type": "string"}}, "required": ["path", "content"]}},
    {"name": "edit_file", "description": "Replace exact text in file.",
     "input_schema": {"type": "object", "properties": {"path": {"type": "string"}, "old_text": {"type": "string"}, "new_text": {"type": "string"}}, "required": ["path", "old_text", "new_text"]}},
    {"name": "spawn_teammate", "description": "Spawn a persistent teammate.",
     "input_schema": {"type": "object", "properties": {"name": {"type": "string"}, "role": {"type": "string"}, "prompt": {"type": "string"}}, "required": ["name", "role", "prompt"]}},
    {"name": "list_teammates", "description": "List all teammates.",
     "input_schema": {"type": "object", "properties": {}}},
    {"name": "send_message", "description": "Send a message to a teammate.",
     "input_schema": {"type": "object", "properties": {"to": {"type": "string"}, "content": {"type": "string"}, "msg_type": {"type": "string", "enum": list(VALID_MSG_TYPES)}}, "required": ["to", "content"]}},
    {"name": "read_inbox", "description": "Read and drain the lead's inbox.",
     "input_schema": {"type": "object", "properties": {}}},
    {"name": "broadcast", "description": "Send a message to all teammates.",
     "input_schema": {"type": "object", "properties": {"content": {"type": "string"}}, "required": ["content"]}},
    {"name": "shutdown_request", "description": "Request a teammate to shut down gracefully. Returns a request_id for tracking.",
     "input_schema": {"type": "object", "properties": {"teammate": {"type": "string"}}, "required": ["teammate"]}},
    {"name": "shutdown_response", "description": "Check the status of a shutdown request by request_id.",
     "input_schema": {"type": "object", "properties": {"request_id": {"type": "string"}}, "required": ["request_id"]}},
    {"name": "plan_approval", "description": "Approve or reject a teammate's plan. Provide request_id + approve + optional feedback.",
     "input_schema": {"type": "object", "properties": {"request_id": {"type": "string"}, "approve": {"type": "boolean"}, "feedback": {"type": "string"}}, "required": ["request_id", "approve"]}},
]


def agent_loop(messages: list):
    while True:
        inbox = BUS.read_inbox("lead")
        if inbox:
            messages.append({
                "role": "user",
                "content": f"<inbox>{json.dumps(inbox, indent=2)}</inbox>",
            })
            messages.append({
                "role": "assistant",
                "content": "Noted inbox messages.",
            })
        response = client.messages.create(
            model=MODEL,
            system=SYSTEM,
            messages=messages,
            tools=TOOLS,
            max_tokens=8000,
        )
        messages.append({"role": "assistant", "content": response.content})
        if response.stop_reason != "tool_use":
            return
        results = []
        for block in response.content:
            if block.type == "tool_use":
                handler = TOOL_HANDLERS.get(block.name)
                try:
                    output = handler(**block.input) if handler else f"Unknown tool: {block.name}"
                except Exception as e:
                    output = f"Error: {e}"
                print(f"> {block.name}: {str(output)[:200]}")
                results.append({
                    "type": "tool_result",
                    "tool_use_id": block.id,
                    "content": str(output),
                })
        messages.append({"role": "user", "content": results})


if __name__ == "__main__":
    history = []
    while True:
        try:
            query = input("\033[36ms10 >> \033[0m")
        except (EOFError, KeyboardInterrupt):
            break
        if query.strip().lower() in ("q", "exit", ""):
            break
        if query.strip() == "/team":
            print(TEAM.list_all())
            continue
        if query.strip() == "/inbox":
            print(json.dumps(BUS.read_inbox("lead"), indent=2))
            continue
        history.append({"role": "user", "content": query})
        agent_loop(history)
        response_content = history[-1]["content"]
        if isinstance(response_content, list):
            for block in response_content:
                if hasattr(block, "text"):
                    print(block.text)
        print()