08_subagents.py

"""08: Programmatic subagents with AgentDefinition, plus SubagentStart/Stop hooks."""

import asyncio

from claude_agent_sdk import (
    AgentDefinition,
    AssistantMessage,
    ClaudeAgentOptions,
    HookMatcher,
    ResultMessage,
    ToolUseBlock,
    query,
)
from demo_util import PROJECT


async def on_subagent_start(input_data, tool_use_id, context):
    print(f"  [hook SubagentStart] type={input_data['agent_type']} id={input_data['agent_id']}")
    return {}


async def on_subagent_stop(input_data, tool_use_id, context):
    print(f"  [hook SubagentStop]  type={input_data['agent_type']} transcript={input_data['agent_transcript_path']}")
    return {}


async def main():
    options = ClaudeAgentOptions(
        cwd=PROJECT,
        allowed_tools=["Read", "Grep", "Glob", "Agent"],
        max_turns=10,
        agents={
            "security-reviewer": AgentDefinition(
                description="Security code reviewer. Use for finding vulnerabilities in source files.",
                prompt="You are a security reviewer. Report each vulnerability as file:line and one sentence. "
                       "Be brief. Never edit files.",
                tools=["Read", "Grep", "Glob"],   # read-only
                model="haiku",                     # cheaper model for the subtask
            ),
        },
        hooks={
            "SubagentStart": [HookMatcher(hooks=[on_subagent_start])],
            "SubagentStop": [HookMatcher(hooks=[on_subagent_stop])],
        },
    )

    async for message in query(
        prompt="Use the security-reviewer agent to review src/app.py, then repeat its findings verbatim.",
        options=options,
    ):
        if isinstance(message, AssistantMessage):
            inside = " (inside subagent)" if message.parent_tool_use_id else ""
            for block in message.content:
                if isinstance(block, ToolUseBlock):
                    if block.name in ("Agent", "Task"):
                        print(f"Subagent invoked: subagent_type={block.input.get('subagent_type')} "
                              f"model={block.input.get('model')}")
                    else:
                        print(f"  tool{inside}: {block.name}")
        elif isinstance(message, ResultMessage):
            print("\nFinal result:\n", message.result)
            print("\nPer-model usage:")
            for model, usage in (message.model_usage or {}).items():
                print(f"  {model}: in={usage['inputTokens']} out={usage['outputTokens']} cost=${usage['costUSD']:.4f}")


asyncio.run(main())