import asyncio

from claude_agent_sdk import ClaudeAgentOptions, query


async def main() -> None:
    async for message in query(
        prompt=(
            "Get a cited adversarial evaluation plan for SupportBot v2.3. "
            "Do not generate adversarial inputs."
        ),
        options=ClaudeAgentOptions(
            mcp_servers={
                "k2-adversarial-context": {
                    "command": "python",
                    "args": ["scripts/k2_adversarial_mcp_server.py"],
                    "env": {
                        "K2_ADV_BUNDLE_DIR": (
                            "docs/customer-demos/demo-adversarial-context/k2-assets"
                        )
                    },
                }
            },
            allowed_tools=["mcp__k2-adversarial-context__get_evaluation_plan"],
        ),
    ):
        print(message)


asyncio.run(main())