16_cost_tracking.py

"""16: Cost and usage: per-step usage deduplicated by message_id, result totals, cache tokens."""

import asyncio

from claude_agent_sdk import AssistantMessage, ClaudeAgentOptions, ResultMessage, query
from demo_util import PROJECT


async def main():
    seen_ids = set()
    input_tokens = 0
    steps = 0

    async for message in query(
        prompt="Read src/utils.py and src/app.py and tell me which one is longer.",
        options=ClaudeAgentOptions(cwd=PROJECT, allowed_tools=["Read"], max_turns=5, effort="low"),
    ):
        if isinstance(message, AssistantMessage) and not message.parent_tool_use_id:
            # Parallel tool calls share one message_id: count each API step once
            if message.message_id not in seen_ids:
                seen_ids.add(message.message_id)
                steps += 1
                u = message.usage or {}
                input_tokens += u.get("input_tokens", 0)
                print(f"step {steps}: message_id={message.message_id[:14]}... "
                      f"input={u.get('input_tokens')} cache_read={u.get('cache_read_input_tokens')} "
                      f"cache_create={u.get('cache_creation_input_tokens')} stop={message.stop_reason}")
        elif isinstance(message, ResultMessage):
            u = message.usage or {}
            print(f"\nresult.usage: input={u.get('input_tokens')} output={u.get('output_tokens')} "
                  f"cache_read={u.get('cache_read_input_tokens')} cache_create={u.get('cache_creation_input_tokens')}")
            print(f"result.total_cost_usd: ${message.total_cost_usd:.4f}  (client-side estimate)")
            print(f"result.duration_ms={message.duration_ms} duration_api_ms={message.duration_api_ms} num_turns={message.num_turns}")
            for model, mu in (message.model_usage or {}).items():
                print(f"result.model_usage[{model}]: cost=${mu['costUSD']:.4f} contextWindow={mu['contextWindow']} "
                      f"costBasis={mu.get('costBasis')}")

    print(f"\nDeduplicated steps: {steps}, summed input tokens: {input_tokens}")


asyncio.run(main())