用 Python SDK 编程造批量 Agent 流水线：Claude Agent SDK 自动化编排实战

pip install claude-agent-sdk

import anyio
from claude_agent_sdk import query

async def main():
    # query() 返回一个异步消息迭代器，逐条 yield Agent 产生的消息
    async for message in query(prompt="What is 2 + 2?"):
        print(message)

anyio.run(main)

python hello.py

from claude_agent_sdk import query, ClaudeAgentOptions, AssistantMessage, TextBlock

options = ClaudeAgentOptions(
    system_prompt="你是资深 Python 安全审计员，只读分析、不修改文件",  # 系统提示词
    max_turns=6,                    # Agent 循环最多迭代几轮，防止失控
    cwd="/home/me/myproject",       # Agent 的工作目录（影响 Read/Bash 的相对路径）
    allowed_tools=["Read", "Bash"], # 白名单：列出的工具自动放行，不再弹确认
    disallowed_tools=["Write", "Edit", "NotebookEdit"],  # 黑名单：彻底移除写操作能力
    permission_mode="acceptEdits",  # 兜底权限模式
)

async for message in query(prompt="列出项目里所有硬编码的密钥", options=options):
    if isinstance(message, AssistantMessage):
        for block in message.content:
            if isinstance(block, TextBlock):
                print(block.text)

import anyio
from claude_agent_sdk import ClaudeAgentOptions, ClaudeSDKClient

async def main():
    options = ClaudeAgentOptions(
        allowed_tools=["Read", "Bash"],
        cwd=".",
    )
    # async with 自动管理 CLI 子进程的生命周期
    async with ClaudeSDKClient(options=options) as client:
        await client.query("读取 main.py 并总结它的职责")
        async for msg in client.receive_response():
            print("第一轮:", msg)

        # 同一会话内继续追问，Agent 记得上文
        await client.query("那它有哪些潜在的安全风险？")
        async for msg in client.receive_response():
            print("第二轮:", msg)

anyio.run(main)

import anyio
from claude_agent_sdk import (
    tool, create_sdk_mcp_server,
    ClaudeAgentOptions, ClaudeSDKClient,
)

# 定义一个工具：@tool(工具名, 描述, 参数类型签名)
@tool("count_secrets", "统计文件中疑似密钥的行数", {"file_path": str})
async def count_secrets(args):
    path = args["file_path"]
    try:
        with open(path, "r", encoding="utf-8") as f:
            lines = f.readlines()
        hits = sum(1 for ln in lines if any(k in ln.upper() for k in ("API_KEY", "SECRET", "TOKEN", "PASSWORD")))
        text = f"{path} 中发现 {hits} 行疑似密钥"
    except FileNotFoundError:
        text = f"文件不存在: {path}"
    # 返回 MCP 标准内容块
    return {"content": [{"type": "text", "text": text}]}

# 把函数打包成进程内 MCP 服务器
server = create_sdk_mcp_server(name="audit-tools", version="1.0.0", tools=[count_secrets])

options = ClaudeAgentOptions(
    mcp_servers={"t": server},
    allowed_tools=["mcp__t__count_secrets"],  # 命名规则：mcp__<服务器名>__<工具名>
)

async def main():
    async with ClaudeSDKClient(options=options) as client:
        await client.query("帮我统计 ./config.py 里有多少疑似密钥，调用 count_secrets 工具")
        async for msg in client.receive_response():
            print(msg)

anyio.run(main)

import anyio
from claude_agent_sdk import ClaudeAgentOptions, ClaudeSDKClient, HookMatcher

async def guard_bash(input_data, tool_use_id, context):
    """PreToolUse 钩子：在 Bash 执行前检查命令"""
    if input_data["tool_name"] != "Bash":
        return {}  # 非 Bash 工具，放行
    command = input_data["tool_input"].get("command", "")
    blocked = ["rm -rf", "DROP TABLE", "sudo ", "> /dev/sda", "mkfs"]
    for pattern in blocked:
        if pattern in command:
            return {
                "hookSpecificOutput": {
                    "hookEventName": "PreToolUse",
                    "permissionDecision": "deny",   # 拒绝执行
                    "permissionDecisionReason": f"命中危险模式: {pattern}",
                }
            }
    return {}  # 放行

options = ClaudeAgentOptions(
    allowed_tools=["Bash", "Read"],
    hooks={
        "PreToolUse": [
            HookMatcher(matcher="Bash", hooks=[guard_bash]),
        ],
    },
)

async def main():
    async with ClaudeSDKClient(options=options) as client:
        await client.query("用 rm -rf 删除 ./tmp 目录")
        async for msg in client.receive_response():
            print(msg)

anyio.run(main)

import asyncio
from claude_agent_sdk import query, ClaudeAgentOptions, AssistantMessage, TextBlock

def make_options():
    """每个子代理用相同的安全配置"""
    return ClaudeAgentOptions(
        system_prompt="你是代码安全审计员。只读分析，输出 JSON：{\"file\":..,\"risks\":[..],\"severity\":\"high|mid|low\"}",
        allowed_tools=["Read", "Bash"],
        disallowed_tools=["Write", "Edit"],
        max_turns=4,
        cwd=".",
    )

async def audit_one(file_path: str) -> dict:
    """一个子代理审一个文件，返回结构化结果"""
    options = make_options()
    collected = []
    prompt = f"审计文件 {file_path}，找出硬编码密钥、SQL 注入、命令注入风险，按约定 JSON 输出。"
    async for message in query(prompt=prompt, options=options):
        if isinstance(message, AssistantMessage):
            for block in message.content:
                if isinstance(block, TextBlock):
                    collected.append(block.text)
    return {"file": file_path, "report": "\n".join(collected)}

async def batch_audit(files: list[str]) -> list[dict]:
    # 并发扇出：所有文件同时审计
    return await asyncio.gather(*[audit_one(f) for f in files])

if __name__ == "__main__":
    targets = ["src/auth.py", "src/db.py", "src/utils.py", "src/api.py"]
    results = asyncio.run(batch_audit(targets))
    for r in results:
        print(f"=== {r['file']} ===")
        print(r["report"])

import asyncio

sem = asyncio.Semaphore(5)  # 最多 5 个子代理并发

async def audit_one(file_path: str) -> dict:
    async with sem:            # 拿到许可才跑
        options = make_options()
        collected = []
        async for message in query(prompt=f"审计 {file_path}", options=options):
            # ... 同上
            pass
        return {"file": file_path, "report": "\n".join(collected)}

import anyio
from claude_agent_sdk import ClaudeAgentOptions, ClaudeSDKClient

async def long_task():
    options = ClaudeAgentOptions(
        allowed_tools=["Read", "Edit", "Bash"],
        cwd=".",
    )
    async with ClaudeSDKClient(options=options) as client:
        # 阶段一：分析
        await client.query("分析整个 src 目录的测试覆盖率缺口")
        async for _ in client.receive_response():
            pass

        # 阶段二：在阶段一的上下文基础上继续
        await client.query("针对最大的三个缺口，各写一个失败的单元测试")
        async for _ in client.receive_response():
            pass
        # 关键：这条 session 的历史在 client 存活期内始终累积

# audit_pipeline.py
import asyncio
import json
from pathlib import Path
from claude_agent_sdk import (
    query, ClaudeAgentOptions, ClaudeSDKClient, HookMatcher,
    AssistantMessage, TextBlock,
)

# ---- 1. 危险命令护栏 ----
async def guard_bash(input_data, tool_use_id, context):
    if input_data["tool_name"] != "Bash":
        return {}
    cmd = input_data["tool_input"].get("command", "")
    for bad in ["rm -rf", "DROP TABLE", "sudo ", ":(){:|:&};:"]:
        if bad in cmd:
            return {"hookSpecificOutput": {
                "hookEventName": "PreToolUse",
                "permissionDecision": "deny",
                "permissionDecisionReason": f"危险命令: {bad}",
            }}
    return {}

# ---- 2. 单文件审计子代理 ----
def audit_options():
    return ClaudeAgentOptions(
        system_prompt=(
            "你是安全审计员。审一个文件，输出严格 JSON："
            '{"file":str,"risks":[{"line":int,"type":str,"desc":str}],'
            '"severity":"high|mid|low|none"}。只读，禁止修改文件。'
        ),
        allowed_tools=["Read", "Bash"],
        disallowed_tools=["Write", "Edit", "NotebookEdit"],
        max_turns=5,
        cwd=".",
        hooks={"PreToolUse": [HookMatcher(matcher="Bash", hooks=[guard_bash])]},
    )

sem = asyncio.Semaphore(4)

async def audit_one(file_path):
    async with sem:
        opts = audit_options()
        text_parts = []
        async for msg in query(
            prompt=f"审计 {file_path} 的安全风险并输出 JSON。",
            options=opts,
        ):
            if isinstance(msg, AssistantMessage):
                for b in msg.content:
                    if isinstance(b, TextBlock):
                        text_parts.append(b.text)
        raw = "\n".join(text_parts)
        # 尝试从回复里抠出 JSON（容错）
        try:
            start, end = raw.index("{"), raw.rindex("}") + 1
            return json.loads(raw[start:end])
        except (ValueError, json.JSONDecodeError):
            return {"file": file_path, "risks": [], "severity": "none",
                    "raw": raw[:200]}

# ---- 3. 编排器 ----
async def run_pipeline(root="."):
    files = [str(p) for p in Path(root).rglob("*.py") if "venv" not in str(p)]
    print(f"发现 {len(files)} 个 Python 文件，开始并行审计（并发上限 4）...")
    results = await asyncio.gather(*[audit_one(f) for f in files])
    high = [r for r in results if r.get("severity") == "high"]
    print(json.dumps({"total": len(files), "high_risk": len(high),
                      "details": high}, ensure_ascii=False, indent=2))

if __name__ == "__main__":
    asyncio.run(run_pipeline("."))

python audit_pipeline.py

flowchart TD
    A[Python 应用] -->|asyncio.run| B[编排器 batch_audit]
    B -->|asyncio.gather 并发扇出| C1[子代理1: query]
    B -->|asyncio.gather 并发扇出| C2[子代理2: query]
    B -->|asyncio.gather 并发扇出| C3[子代理N: query]
    B -.->|Semaphore=4 限流| C1

    C1 --> D1[ClaudeAgentOptions]
    D1 --> D1a[allowed_tools 白名单]
    D1 --> D1b[disallowed_tools 黑名单]
    D1 --> D1c[hooks PreToolUse]

    C1 -->|Agent 循环| E1[模型决策]
    E1 -->|调用工具前| F1[Hook guard_bash 拦截]
    F1 -->|deny 危险命令| G1[拒绝执行]
    F1 -->|allow 安全命令| H1[执行 Read/Bash]
    H1 --> E1
    E1 -->|完成| I1[结构化 JSON 结果]
    C2 --> I2[结果]
    C3 --> I3[结果]

    I1 --> J[汇总报告]
    I2 --> J
    I3 --> J
    J --> K[CI / PR 评论 / 告警]

    style F1 fill:#fde68a,stroke:#b45309
    style G1 fill:#fecaca,stroke:#b91c1c
    style B fill:#bfdbfe,stroke:#1d4ed8