diff --git a/pyproject.toml b/pyproject.toml index 2038ccd81f..2c6ae45eb4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -83,6 +83,10 @@ ignore_missing_imports = true module = "anthropic.*" ignore_missing_imports = true +[[tool.mypy.overrides]] +module = "claude_agent_sdk.*" +ignore_missing_imports = true + [[tool.mypy.overrides]] module = "sanic.*" ignore_missing_imports = true diff --git a/scripts/populate_tox/config.py b/scripts/populate_tox/config.py index 37d3a6a64d..781a590a74 100644 --- a/scripts/populate_tox/config.py +++ b/scripts/populate_tox/config.py @@ -78,6 +78,13 @@ }, "num_versions": 2, }, + "claude_agent_sdk": { + "package": "claude-agent-sdk", + "deps": { + "*": ["pytest-asyncio"], + }, + "python": ">=3.10", + }, "clickhouse_driver": { "package": "clickhouse-driver", "num_versions": 2, diff --git a/scripts/split_tox_gh_actions/split_tox_gh_actions.py b/scripts/split_tox_gh_actions/split_tox_gh_actions.py index b59e768a56..d62ac6450b 100755 --- a/scripts/split_tox_gh_actions/split_tox_gh_actions.py +++ b/scripts/split_tox_gh_actions/split_tox_gh_actions.py @@ -74,6 +74,7 @@ "fastmcp", ], "Agents": [ + "claude_agent_sdk", "openai_agents", "pydantic_ai", ], diff --git a/sentry_sdk/integrations/__init__.py b/sentry_sdk/integrations/__init__.py index 9c76dfe471..c81b8fa3a8 100644 --- a/sentry_sdk/integrations/__init__.py +++ b/sentry_sdk/integrations/__init__.py @@ -67,6 +67,7 @@ def iter_default_integrations( _AUTO_ENABLING_INTEGRATIONS = [ "sentry_sdk.integrations.aiohttp.AioHttpIntegration", "sentry_sdk.integrations.anthropic.AnthropicIntegration", + "sentry_sdk.integrations.claude_agent_sdk.ClaudeAgentSDKIntegration", "sentry_sdk.integrations.ariadne.AriadneIntegration", "sentry_sdk.integrations.arq.ArqIntegration", "sentry_sdk.integrations.asyncpg.AsyncPGIntegration", @@ -127,6 +128,7 @@ def iter_default_integrations( "celery": (4, 4, 7), "chalice": (1, 16, 0), "clickhouse_driver": (0, 2, 0), + "claude_agent_sdk": (0, 1, 0), "cohere": (5, 4, 0), "django": (1, 8), "dramatiq": (1, 9), diff --git a/sentry_sdk/integrations/claude_agent_sdk.py b/sentry_sdk/integrations/claude_agent_sdk.py new file mode 100644 index 0000000000..604c8d617d --- /dev/null +++ b/sentry_sdk/integrations/claude_agent_sdk.py @@ -0,0 +1,488 @@ +from functools import wraps +from typing import TYPE_CHECKING + +import sentry_sdk +from sentry_sdk.ai.monitoring import record_token_usage +from sentry_sdk.ai.utils import set_data_normalized, get_start_span_function +from sentry_sdk.consts import OP, SPANDATA, SPANSTATUS +from sentry_sdk.integrations import _check_minimum_version, DidNotEnable, Integration +from sentry_sdk.scope import should_send_default_pii +from sentry_sdk.utils import ( + capture_internal_exceptions, + event_from_exception, + package_version, +) +from sentry_sdk.tracing_utils import set_span_errored + +try: + import claude_agent_sdk + from claude_agent_sdk import ( + query as original_query, + ClaudeSDKClient, + ) +except ImportError: + raise DidNotEnable("claude-agent-sdk not installed") + +if TYPE_CHECKING: + from typing import Any, AsyncGenerator, Optional + from sentry_sdk.tracing import Span + +AGENT_NAME = "claude-agent" +GEN_AI_SYSTEM = "claude-agent-sdk-python" + + +def _is_assistant_message(message: "Any") -> bool: + """Check if message is an AssistantMessage using duck typing.""" + return hasattr(message, "content") and hasattr(message, "model") + + +def _is_result_message(message: "Any") -> bool: + """Check if message is a ResultMessage using duck typing.""" + return hasattr(message, "usage") and hasattr(message, "total_cost_usd") + + +def _is_text_block(block: "Any") -> bool: + """Check if block is a TextBlock using duck typing.""" + # TextBlock has 'text' but not 'tool_use_id' or 'name' + return ( + hasattr(block, "text") + and not hasattr(block, "tool_use_id") + and not hasattr(block, "name") + ) + + +def _is_tool_use_block(block: "Any") -> bool: + """Check if block is a ToolUseBlock using duck typing.""" + # ToolUseBlock has 'id', 'name', and 'input' + return hasattr(block, "id") and hasattr(block, "name") and hasattr(block, "input") + + +def _is_tool_result_block(block: "Any") -> bool: + """Check if block is a ToolResultBlock using duck typing.""" + # ToolResultBlock has 'tool_use_id' + return hasattr(block, "tool_use_id") + + +class ClaudeAgentSDKIntegration(Integration): + identifier = "claude_agent_sdk" + origin = f"auto.ai.{identifier}" + + def __init__(self, include_prompts: bool = True) -> None: + self.include_prompts = include_prompts + + @staticmethod + def setup_once() -> None: + version = package_version("claude_agent_sdk") + _check_minimum_version(ClaudeAgentSDKIntegration, version) + claude_agent_sdk.query = _wrap_query(original_query) + ClaudeSDKClient.query = _wrap_client_query(ClaudeSDKClient.query) + ClaudeSDKClient.receive_response = _wrap_receive_response( + ClaudeSDKClient.receive_response + ) + + +def _should_include_prompts(integration: "ClaudeAgentSDKIntegration") -> bool: + return should_send_default_pii() and integration.include_prompts + + +def _capture_exception(exc: "Any") -> None: + set_span_errored() + event, hint = event_from_exception( + exc, + client_options=sentry_sdk.get_client().options, + mechanism={"type": "claude_agent_sdk", "handled": False}, + ) + sentry_sdk.capture_event(event, hint=hint) + + +def _set_span_input_data( + span: "Span", + prompt: str, + options: "Optional[Any]", + integration: "ClaudeAgentSDKIntegration", +) -> None: + set_data_normalized(span, SPANDATA.GEN_AI_SYSTEM, GEN_AI_SYSTEM) + set_data_normalized(span, SPANDATA.GEN_AI_OPERATION_NAME, "chat") + + if options is not None: + model = getattr(options, "model", None) + if model: + set_data_normalized(span, SPANDATA.GEN_AI_REQUEST_MODEL, model) + + allowed_tools = getattr(options, "allowed_tools", None) + if allowed_tools: + tools_list = [{"name": tool} for tool in allowed_tools] + set_data_normalized( + span, SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS, tools_list, unpack=False + ) + + if _should_include_prompts(integration): + messages = [] + system_prompt = getattr(options, "system_prompt", None) if options else None + if system_prompt: + messages.append({"role": "system", "content": system_prompt}) + messages.append({"role": "user", "content": prompt}) + set_data_normalized( + span, SPANDATA.GEN_AI_REQUEST_MESSAGES, messages, unpack=False + ) + + +def _extract_text_from_message(message: "Any") -> "Optional[str]": + if not _is_assistant_message(message): + return None + text_parts = [ + block.text for block in getattr(message, "content", []) if _is_text_block(block) + ] + return "".join(text_parts) if text_parts else None + + +def _extract_tool_calls(message: "Any") -> "Optional[list]": + if not _is_assistant_message(message): + return None + tool_calls = [] + for block in getattr(message, "content", []): + if _is_tool_use_block(block): + tool_call = {"name": getattr(block, "name", "unknown")} + tool_input = getattr(block, "input", None) + if tool_input is not None: + tool_call["input"] = tool_input + tool_calls.append(tool_call) + return tool_calls or None + + +def _extract_message_data(messages: list) -> dict: + """Extract relevant data from a list of messages.""" + data = { + "response_texts": [], + "tool_calls": [], + "total_cost": None, + "input_tokens": None, + "output_tokens": None, + "cached_input_tokens": None, + "response_model": None, + } + + for message in messages: + if _is_assistant_message(message): + text = _extract_text_from_message(message) + if text: + data["response_texts"].append(text) + + calls = _extract_tool_calls(message) + if calls: + data["tool_calls"].extend(calls) + + if not data["response_model"]: + data["response_model"] = getattr(message, "model", None) + + elif _is_result_message(message): + data["total_cost"] = getattr(message, "total_cost_usd", None) + usage = getattr(message, "usage", None) + if isinstance(usage, dict): + data["input_tokens"] = usage.get("input_tokens") + data["output_tokens"] = usage.get("output_tokens") + # Store cached tokens separately for the backend to apply discount pricing + cached_input = usage.get("cache_read_input_tokens") or 0 + data["cached_input_tokens"] = cached_input if cached_input > 0 else None + + return data + + +def _set_span_output_data( + span: "Span", + messages: list, + integration: "ClaudeAgentSDKIntegration", +) -> None: + data = _extract_message_data(messages) + + if data["response_model"]: + set_data_normalized( + span, SPANDATA.GEN_AI_RESPONSE_MODEL, data["response_model"] + ) + if SPANDATA.GEN_AI_REQUEST_MODEL not in getattr(span, "_data", {}): + set_data_normalized( + span, SPANDATA.GEN_AI_REQUEST_MODEL, data["response_model"] + ) + + if _should_include_prompts(integration): + if data["response_texts"]: + set_data_normalized( + span, SPANDATA.GEN_AI_RESPONSE_TEXT, data["response_texts"] + ) + if data["tool_calls"]: + set_data_normalized( + span, + SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS, + data["tool_calls"], + unpack=False, + ) + + if data["input_tokens"] is not None or data["output_tokens"] is not None: + record_token_usage( + span, + input_tokens=data["input_tokens"], + input_tokens_cached=data["cached_input_tokens"], + output_tokens=data["output_tokens"], + ) + + if data["total_cost"] is not None: + span.set_data("claude_code.total_cost_usd", data["total_cost"]) + + +def _start_invoke_agent_span( + prompt: str, + options: "Optional[Any]", + integration: "ClaudeAgentSDKIntegration", +) -> "Span": + span = get_start_span_function()( + op=OP.GEN_AI_INVOKE_AGENT, + name=f"invoke_agent {AGENT_NAME}", + origin=ClaudeAgentSDKIntegration.origin, + ) + span.__enter__() + + set_data_normalized(span, SPANDATA.GEN_AI_OPERATION_NAME, "invoke_agent") + set_data_normalized(span, SPANDATA.GEN_AI_AGENT_NAME, AGENT_NAME) + set_data_normalized(span, SPANDATA.GEN_AI_SYSTEM, GEN_AI_SYSTEM) + + if _should_include_prompts(integration): + messages = [] + system_prompt = getattr(options, "system_prompt", None) if options else None + if system_prompt: + messages.append({"role": "system", "content": system_prompt}) + messages.append({"role": "user", "content": prompt}) + set_data_normalized( + span, SPANDATA.GEN_AI_REQUEST_MESSAGES, messages, unpack=False + ) + + return span + + +def _end_invoke_agent_span( + span: "Span", + messages: list, + integration: "ClaudeAgentSDKIntegration", +) -> None: + data = _extract_message_data(messages) + + if _should_include_prompts(integration) and data["response_texts"]: + set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, data["response_texts"]) + + if data["response_model"]: + set_data_normalized( + span, SPANDATA.GEN_AI_RESPONSE_MODEL, data["response_model"] + ) + set_data_normalized(span, SPANDATA.GEN_AI_REQUEST_MODEL, data["response_model"]) + + if data["input_tokens"] is not None or data["output_tokens"] is not None: + record_token_usage( + span, + input_tokens=data["input_tokens"], + input_tokens_cached=data["cached_input_tokens"], + output_tokens=data["output_tokens"], + ) + + if data["total_cost"] is not None: + span.set_data("claude_code.total_cost_usd", data["total_cost"]) + + span.__exit__(None, None, None) + + +def _create_execute_tool_span( + tool_use: "Any", + tool_result: "Optional[Any]", + integration: "ClaudeAgentSDKIntegration", +) -> "Span": + tool_name = getattr(tool_use, "name", "unknown") + span = sentry_sdk.start_span( + op=OP.GEN_AI_EXECUTE_TOOL, + name=f"execute_tool {tool_name}", + origin=ClaudeAgentSDKIntegration.origin, + ) + + set_data_normalized(span, SPANDATA.GEN_AI_OPERATION_NAME, "execute_tool") + set_data_normalized(span, SPANDATA.GEN_AI_TOOL_NAME, tool_name) + set_data_normalized(span, SPANDATA.GEN_AI_SYSTEM, GEN_AI_SYSTEM) + + if _should_include_prompts(integration): + tool_input = getattr(tool_use, "input", None) + if tool_input is not None: + set_data_normalized(span, SPANDATA.GEN_AI_TOOL_INPUT, tool_input) + + if tool_result is not None: + tool_output = getattr(tool_result, "content", None) + if tool_output is not None: + set_data_normalized(span, SPANDATA.GEN_AI_TOOL_OUTPUT, tool_output) + + if tool_result is not None and getattr(tool_result, "is_error", False): + span.set_status(SPANSTATUS.INTERNAL_ERROR) + + return span + + +def _process_tool_executions( + messages: list, integration: "ClaudeAgentSDKIntegration" +) -> list: + """Create execute_tool spans for tool executions found in messages. + + Returns a list of the created spans (for testing purposes). + """ + tool_uses = {} + tool_results = {} + + for message in messages: + if not _is_assistant_message(message): + continue + for block in getattr(message, "content", []): + if _is_tool_use_block(block): + tool_id = getattr(block, "id", None) + if tool_id: + tool_uses[tool_id] = block + elif _is_tool_result_block(block): + tool_use_id = getattr(block, "tool_use_id", None) + if tool_use_id: + tool_results[tool_use_id] = block + + spans = [] + for tool_id, tool_use in tool_uses.items(): + span = _create_execute_tool_span( + tool_use, tool_results.get(tool_id), integration + ) + span.finish() + spans.append(span) + return spans + + +def _wrap_query(original_func: "Any") -> "Any": + @wraps(original_func) + async def wrapper( + *, prompt: str, options: "Optional[Any]" = None, **kwargs: "Any" + ) -> "AsyncGenerator[Any, None]": + integration = sentry_sdk.get_client().get_integration(ClaudeAgentSDKIntegration) + if integration is None: + async for message in original_func( + prompt=prompt, options=options, **kwargs + ): + yield message + return + + model = getattr(options, "model", "") if options else "" + invoke_span = _start_invoke_agent_span(prompt, options, integration) + + chat_span = get_start_span_function()( + op=OP.GEN_AI_CHAT, + name=f"claude-agent-sdk query {model}".strip(), + origin=ClaudeAgentSDKIntegration.origin, + ) + chat_span.__enter__() + + with capture_internal_exceptions(): + _set_span_input_data(chat_span, prompt, options, integration) + + collected_messages = [] + try: + async for message in original_func( + prompt=prompt, options=options, **kwargs + ): + collected_messages.append(message) + yield message + except Exception as exc: + _capture_exception(exc) + raise + finally: + with capture_internal_exceptions(): + _set_span_output_data(chat_span, collected_messages, integration) + chat_span.__exit__(None, None, None) + + with capture_internal_exceptions(): + _process_tool_executions(collected_messages, integration) + + with capture_internal_exceptions(): + _end_invoke_agent_span(invoke_span, collected_messages, integration) + + return wrapper + + +def _wrap_client_query(original_method: "Any") -> "Any": + @wraps(original_method) + async def wrapper(self: "Any", prompt: str, **kwargs: "Any") -> "Any": + integration = sentry_sdk.get_client().get_integration(ClaudeAgentSDKIntegration) + if integration is None: + return await original_method(self, prompt, **kwargs) + + options = getattr(self, "_options", None) + model = getattr(options, "model", "") if options else "" + + invoke_span = _start_invoke_agent_span(prompt, options, integration) + + chat_span = get_start_span_function()( + op=OP.GEN_AI_CHAT, + name=f"claude-agent-sdk client {model}".strip(), + origin=ClaudeAgentSDKIntegration.origin, + ) + chat_span.__enter__() + + with capture_internal_exceptions(): + _set_span_input_data(chat_span, prompt, options, integration) + + self._sentry_query_context = { + "invoke_span": invoke_span, + "chat_span": chat_span, + "integration": integration, + "messages": [], + } + + try: + return await original_method(self, prompt, **kwargs) + except Exception as exc: + _capture_exception(exc) + messages = self._sentry_query_context.get("messages", []) + with capture_internal_exceptions(): + _set_span_output_data(chat_span, messages, integration) + chat_span.__exit__(None, None, None) + with capture_internal_exceptions(): + _end_invoke_agent_span(invoke_span, messages, integration) + self._sentry_query_context = {} + raise + + return wrapper + + +def _wrap_receive_response(original_method: "Any") -> "Any": + @wraps(original_method) + async def wrapper(self: "Any", **kwargs: "Any") -> "AsyncGenerator[Any, None]": + integration = sentry_sdk.get_client().get_integration(ClaudeAgentSDKIntegration) + if integration is None: + async for message in original_method(self, **kwargs): + yield message + return + + context = getattr(self, "_sentry_query_context", {}) + invoke_span = context.get("invoke_span") + chat_span = context.get("chat_span") + stored_integration = context.get("integration", integration) + messages = context.get("messages", []) + + try: + async for message in original_method(self, **kwargs): + messages.append(message) + yield message + except Exception as exc: + _capture_exception(exc) + raise + finally: + if chat_span is not None: + with capture_internal_exceptions(): + _set_span_output_data(chat_span, messages, stored_integration) + chat_span.__exit__(None, None, None) + + with capture_internal_exceptions(): + _process_tool_executions(messages, stored_integration) + + if invoke_span is not None: + with capture_internal_exceptions(): + _end_invoke_agent_span(invoke_span, messages, stored_integration) + + self._sentry_query_context = {} + + return wrapper diff --git a/setup.py b/setup.py index be8e82b26f..d36f3dc772 100644 --- a/setup.py +++ b/setup.py @@ -52,6 +52,7 @@ def get_file_text(file_name): "celery": ["celery>=3"], "celery-redbeat": ["celery-redbeat>=2"], "chalice": ["chalice>=1.16.0"], + "claude_agent_sdk": ["claude-agent-sdk>=0.1.0"], "clickhouse-driver": ["clickhouse-driver>=0.2.0"], "django": ["django>=1.8"], "falcon": ["falcon>=1.4"], diff --git a/tests/integrations/claude_agent_sdk/__init__.py b/tests/integrations/claude_agent_sdk/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/integrations/claude_agent_sdk/test_claude_agent_sdk.py b/tests/integrations/claude_agent_sdk/test_claude_agent_sdk.py new file mode 100644 index 0000000000..973584a04b --- /dev/null +++ b/tests/integrations/claude_agent_sdk/test_claude_agent_sdk.py @@ -0,0 +1,853 @@ +import pytest +from dataclasses import dataclass +from typing import List, Optional +import json + +from sentry_sdk import start_transaction +from sentry_sdk.consts import OP, SPANDATA +from sentry_sdk.integrations.claude_agent_sdk import ( + ClaudeAgentSDKIntegration, + _set_span_input_data, + _set_span_output_data, + _extract_text_from_message, + _extract_tool_calls, + _start_invoke_agent_span, + _end_invoke_agent_span, + _create_execute_tool_span, + _process_tool_executions, + _wrap_query, + AGENT_NAME, +) +from claude_agent_sdk import ( + AssistantMessage, + ResultMessage, + TextBlock, + ToolUseBlock, + ToolResultBlock, +) + + +@dataclass +class Options: + model: Optional[str] = None + allowed_tools: Optional[List[str]] = None + system_prompt: Optional[str] = None + + +def make_result_message(usage=None, total_cost_usd=None): + return ResultMessage( + subtype="result", + duration_ms=1000, + duration_api_ms=900, + is_error=False, + num_turns=1, + session_id="test-session", + total_cost_usd=total_cost_usd, + usage=usage, + ) + + +@pytest.fixture +def integration(): + return ClaudeAgentSDKIntegration(include_prompts=True) + + +@pytest.fixture +def integration_no_prompts(): + return ClaudeAgentSDKIntegration(include_prompts=False) + + +@pytest.fixture +def assistant_message(): + return AssistantMessage( + content=[TextBlock(text="Hello! I'm Claude.")], + model="claude-sonnet-4-5-20250929", + ) + + +@pytest.fixture +def result_message(): + return make_result_message( + usage={ + "input_tokens": 10, + "output_tokens": 20, + "cache_read_input_tokens": 100, + }, + total_cost_usd=0.005, + ) + + +def test_extract_text_from_assistant_message(): + message = AssistantMessage( + content=[TextBlock(text="Hello!")], + model="test-model", + ) + assert _extract_text_from_message(message) == "Hello!" + + +def test_extract_text_from_multiple_blocks(): + message = AssistantMessage( + content=[ + TextBlock(text="First. "), + TextBlock(text="Second."), + ], + model="test-model", + ) + assert _extract_text_from_message(message) == "First. Second." + + +def test_extract_text_returns_none_for_non_assistant(): + result = make_result_message(usage=None) + assert _extract_text_from_message(result) is None + + +def test_extract_tool_calls(): + message = AssistantMessage( + content=[ + TextBlock(text="Let me help."), + ToolUseBlock(id="tool-1", name="Read", input={"path": "/test.txt"}), + ], + model="test-model", + ) + tool_calls = _extract_tool_calls(message) + assert len(tool_calls) == 1 + assert tool_calls[0]["name"] == "Read" + assert tool_calls[0]["input"] == {"path": "/test.txt"} + + +def test_extract_tool_calls_returns_none_for_non_assistant(): + result = make_result_message(usage=None) + assert _extract_tool_calls(result) is None + + +def test_set_span_input_data_basic(sentry_init, integration): + sentry_init( + integrations=[ClaudeAgentSDKIntegration()], + traces_sample_rate=1.0, + send_default_pii=True, + ) + + with start_transaction(name="test") as transaction: + span = transaction.start_child(op="test") + _set_span_input_data(span, "Hello", None, integration) + + assert span._data[SPANDATA.GEN_AI_SYSTEM] == "claude-agent-sdk-python" + assert span._data[SPANDATA.GEN_AI_OPERATION_NAME] == "chat" + assert SPANDATA.GEN_AI_REQUEST_MESSAGES in span._data + + +def test_set_span_input_data_with_options(sentry_init, integration): + sentry_init( + integrations=[ClaudeAgentSDKIntegration()], + traces_sample_rate=1.0, + send_default_pii=True, + ) + + with start_transaction(name="test") as transaction: + span = transaction.start_child(op="test") + options = Options( + model="claude-opus-4-5-20251101", + allowed_tools=["Read", "Write"], + system_prompt="You are helpful.", + ) + _set_span_input_data(span, "Hello", options, integration) + + assert span._data[SPANDATA.GEN_AI_REQUEST_MODEL] == "claude-opus-4-5-20251101" + assert SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS in span._data + messages = json.loads(span._data[SPANDATA.GEN_AI_REQUEST_MESSAGES]) + assert len(messages) == 2 + assert messages[0]["role"] == "system" + assert messages[0]["content"] == "You are helpful." + assert messages[1]["role"] == "user" + + +def test_set_span_input_data_pii_disabled(sentry_init, integration): + sentry_init( + integrations=[ClaudeAgentSDKIntegration()], + traces_sample_rate=1.0, + send_default_pii=False, + ) + + with start_transaction(name="test") as transaction: + span = transaction.start_child(op="test") + _set_span_input_data(span, "Hello", None, integration) + + assert span._data[SPANDATA.GEN_AI_SYSTEM] == "claude-agent-sdk-python" + assert SPANDATA.GEN_AI_REQUEST_MESSAGES not in span._data + + +def test_set_span_input_data_include_prompts_disabled( + sentry_init, integration_no_prompts +): + sentry_init( + integrations=[ClaudeAgentSDKIntegration(include_prompts=False)], + traces_sample_rate=1.0, + send_default_pii=True, + ) + + with start_transaction(name="test") as transaction: + span = transaction.start_child(op="test") + _set_span_input_data(span, "Hello", None, integration_no_prompts) + + assert span._data[SPANDATA.GEN_AI_SYSTEM] == "claude-agent-sdk-python" + assert SPANDATA.GEN_AI_REQUEST_MESSAGES not in span._data + + +def test_set_span_output_data_with_messages( + sentry_init, integration, assistant_message, result_message +): + sentry_init( + integrations=[ClaudeAgentSDKIntegration()], + traces_sample_rate=1.0, + send_default_pii=True, + ) + + with start_transaction(name="test") as transaction: + span = transaction.start_child(op="test") + messages = [assistant_message, result_message] + _set_span_output_data(span, messages, integration) + + assert span._data[SPANDATA.GEN_AI_RESPONSE_MODEL] == "claude-sonnet-4-5-20250929" + assert span._data[SPANDATA.GEN_AI_REQUEST_MODEL] == "claude-sonnet-4-5-20250929" + assert span._data[SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 10 + assert span._data[SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] == 20 + assert span._data[SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS] == 30 + assert span._data[SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHED] == 100 + assert span._data["claude_code.total_cost_usd"] == 0.005 + + +def test_set_span_output_data_no_usage(sentry_init, integration, assistant_message): + sentry_init( + integrations=[ClaudeAgentSDKIntegration()], + traces_sample_rate=1.0, + send_default_pii=True, + ) + + with start_transaction(name="test") as transaction: + span = transaction.start_child(op="test") + result_no_usage = make_result_message(usage=None, total_cost_usd=None) + messages = [assistant_message, result_no_usage] + _set_span_output_data(span, messages, integration) + + assert span._data[SPANDATA.GEN_AI_RESPONSE_MODEL] == "claude-sonnet-4-5-20250929" + assert SPANDATA.GEN_AI_USAGE_INPUT_TOKENS not in span._data + assert "claude_code.total_cost_usd" not in span._data + + +def test_set_span_output_data_with_tool_calls(sentry_init, integration, result_message): + sentry_init( + integrations=[ClaudeAgentSDKIntegration()], + traces_sample_rate=1.0, + send_default_pii=True, + ) + + with start_transaction(name="test") as transaction: + span = transaction.start_child(op="test") + assistant_with_tool = AssistantMessage( + content=[ + TextBlock(text="Let me read that."), + ToolUseBlock(id="tool-1", name="Read", input={"path": "/test.txt"}), + ], + model="claude-sonnet-4-5-20250929", + ) + messages = [assistant_with_tool, result_message] + _set_span_output_data(span, messages, integration) + + assert SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS in span._data + + +def test_set_span_output_data_pii_disabled( + sentry_init, integration, assistant_message, result_message +): + sentry_init( + integrations=[ClaudeAgentSDKIntegration()], + traces_sample_rate=1.0, + send_default_pii=False, + ) + + with start_transaction(name="test") as transaction: + span = transaction.start_child(op="test") + messages = [assistant_message, result_message] + _set_span_output_data(span, messages, integration) + + assert span._data[SPANDATA.GEN_AI_RESPONSE_MODEL] == "claude-sonnet-4-5-20250929" + assert span._data[SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 10 + assert SPANDATA.GEN_AI_RESPONSE_TEXT not in span._data + + +def test_empty_messages_list(sentry_init, integration): + sentry_init( + integrations=[ClaudeAgentSDKIntegration()], + traces_sample_rate=1.0, + send_default_pii=True, + ) + + with start_transaction(name="test") as transaction: + span = transaction.start_child(op="test") + _set_span_output_data(span, [], integration) + + assert SPANDATA.GEN_AI_RESPONSE_MODEL not in span._data + assert SPANDATA.GEN_AI_RESPONSE_TEXT not in span._data + + +def test_integration_identifier(): + integration = ClaudeAgentSDKIntegration() + assert integration.identifier == "claude_agent_sdk" + assert integration.origin == "auto.ai.claude_agent_sdk" + + +def test_integration_include_prompts_default(): + integration = ClaudeAgentSDKIntegration() + assert integration.include_prompts is True + + +def test_integration_include_prompts_false(): + integration = ClaudeAgentSDKIntegration(include_prompts=False) + assert integration.include_prompts is False + + +@pytest.mark.parametrize( + "send_default_pii,include_prompts,expect_messages", + [ + (True, True, True), + (True, False, False), + (False, True, False), + (False, False, False), + ], +) +def test_pii_and_prompts_matrix( + sentry_init, send_default_pii, include_prompts, expect_messages +): + sentry_init( + integrations=[ClaudeAgentSDKIntegration(include_prompts=include_prompts)], + traces_sample_rate=1.0, + send_default_pii=send_default_pii, + ) + + with start_transaction(name="test") as transaction: + span = transaction.start_child(op="test") + integration = ClaudeAgentSDKIntegration(include_prompts=include_prompts) + _set_span_input_data(span, "Test prompt", None, integration) + + if expect_messages: + assert SPANDATA.GEN_AI_REQUEST_MESSAGES in span._data + else: + assert SPANDATA.GEN_AI_REQUEST_MESSAGES not in span._data + + +def test_model_fallback_from_response( + sentry_init, integration, assistant_message, result_message +): + sentry_init( + integrations=[ClaudeAgentSDKIntegration()], + traces_sample_rate=1.0, + send_default_pii=True, + ) + + with start_transaction(name="test") as transaction: + span = transaction.start_child(op="test") + _set_span_input_data(span, "Hello", None, integration) + messages = [assistant_message, result_message] + _set_span_output_data(span, messages, integration) + + assert span._data[SPANDATA.GEN_AI_REQUEST_MODEL] == "claude-sonnet-4-5-20250929" + assert span._data[SPANDATA.GEN_AI_RESPONSE_MODEL] == "claude-sonnet-4-5-20250929" + + +def test_model_from_options_preserved( + sentry_init, integration, assistant_message, result_message +): + sentry_init( + integrations=[ClaudeAgentSDKIntegration()], + traces_sample_rate=1.0, + send_default_pii=True, + ) + + with start_transaction(name="test") as transaction: + span = transaction.start_child(op="test") + options = Options(model="claude-opus-4-5-20251101") + _set_span_input_data(span, "Hello", options, integration) + messages = [assistant_message, result_message] + _set_span_output_data(span, messages, integration) + + assert span._data[SPANDATA.GEN_AI_REQUEST_MODEL] == "claude-opus-4-5-20251101" + assert span._data[SPANDATA.GEN_AI_RESPONSE_MODEL] == "claude-sonnet-4-5-20250929" + + +def test_available_tools_format(sentry_init, integration): + sentry_init( + integrations=[ClaudeAgentSDKIntegration()], + traces_sample_rate=1.0, + send_default_pii=True, + ) + + with start_transaction(name="test") as transaction: + span = transaction.start_child(op="test") + options = Options(allowed_tools=["Read", "Write", "Bash"]) + _set_span_input_data(span, "Hello", options, integration) + + tools_data = span._data[SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS] + assert isinstance(tools_data, str) + tools = json.loads(tools_data) + assert len(tools) == 3 + assert {"name": "Read"} in tools + assert {"name": "Write"} in tools + assert {"name": "Bash"} in tools + + +def test_cached_tokens_extraction(sentry_init, integration, assistant_message): + sentry_init( + integrations=[ClaudeAgentSDKIntegration()], + traces_sample_rate=1.0, + send_default_pii=True, + ) + + with start_transaction(name="test") as transaction: + span = transaction.start_child(op="test") + result_with_cache = make_result_message( + usage={ + "input_tokens": 5, + "output_tokens": 15, + "cache_read_input_tokens": 500, + }, + total_cost_usd=0.003, + ) + messages = [assistant_message, result_with_cache] + _set_span_output_data(span, messages, integration) + + assert span._data[SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 5 + assert span._data[SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] == 15 + assert span._data[SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS] == 20 + assert span._data[SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHED] == 500 + + +def test_start_invoke_agent_span_basic(sentry_init, integration): + sentry_init( + integrations=[ClaudeAgentSDKIntegration()], + traces_sample_rate=1.0, + send_default_pii=True, + ) + + with start_transaction(name="test"): + span = _start_invoke_agent_span("Hello", None, integration) + span.__exit__(None, None, None) + + assert span.op == OP.GEN_AI_INVOKE_AGENT + assert span.description == f"invoke_agent {AGENT_NAME}" + assert span._data[SPANDATA.GEN_AI_OPERATION_NAME] == "invoke_agent" + assert span._data[SPANDATA.GEN_AI_AGENT_NAME] == AGENT_NAME + assert span._data[SPANDATA.GEN_AI_SYSTEM] == "claude-agent-sdk-python" + assert SPANDATA.GEN_AI_REQUEST_MESSAGES in span._data + + +def test_start_invoke_agent_span_with_system_prompt(sentry_init, integration): + sentry_init( + integrations=[ClaudeAgentSDKIntegration()], + traces_sample_rate=1.0, + send_default_pii=True, + ) + + with start_transaction(name="test"): + options = Options(system_prompt="You are helpful.") + span = _start_invoke_agent_span("Hello", options, integration) + span.__exit__(None, None, None) + + messages = json.loads(span._data[SPANDATA.GEN_AI_REQUEST_MESSAGES]) + assert len(messages) == 2 + assert messages[0]["role"] == "system" + assert messages[0]["content"] == "You are helpful." + assert messages[1]["role"] == "user" + assert messages[1]["content"] == "Hello" + + +def test_start_invoke_agent_span_pii_disabled(sentry_init, integration): + sentry_init( + integrations=[ClaudeAgentSDKIntegration()], + traces_sample_rate=1.0, + send_default_pii=False, + ) + + with start_transaction(name="test"): + span = _start_invoke_agent_span("Hello", None, integration) + span.__exit__(None, None, None) + + assert span._data[SPANDATA.GEN_AI_SYSTEM] == "claude-agent-sdk-python" + assert SPANDATA.GEN_AI_REQUEST_MESSAGES not in span._data + + +def test_end_invoke_agent_span_aggregates_data( + sentry_init, integration, assistant_message, result_message +): + sentry_init( + integrations=[ClaudeAgentSDKIntegration()], + traces_sample_rate=1.0, + send_default_pii=True, + ) + + with start_transaction(name="test"): + span = _start_invoke_agent_span("Hello", None, integration) + messages = [assistant_message, result_message] + _end_invoke_agent_span(span, messages, integration) + + assert span._data[SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 10 + assert span._data[SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] == 20 + assert span._data[SPANDATA.GEN_AI_RESPONSE_MODEL] == "claude-sonnet-4-5-20250929" + + +def test_create_execute_tool_span_basic(sentry_init, integration): + sentry_init( + integrations=[ClaudeAgentSDKIntegration()], + traces_sample_rate=1.0, + send_default_pii=True, + ) + + with start_transaction(name="test"): + tool_use = ToolUseBlock(id="tool-1", name="Read", input={"path": "/test.txt"}) + span = _create_execute_tool_span(tool_use, None, integration) + span.finish() + + assert span.op == OP.GEN_AI_EXECUTE_TOOL + assert span.description == "execute_tool Read" + assert span._data[SPANDATA.GEN_AI_OPERATION_NAME] == "execute_tool" + assert span._data[SPANDATA.GEN_AI_TOOL_NAME] == "Read" + assert span._data[SPANDATA.GEN_AI_SYSTEM] == "claude-agent-sdk-python" + + +def test_create_execute_tool_span_with_result(sentry_init, integration): + sentry_init( + integrations=[ClaudeAgentSDKIntegration()], + traces_sample_rate=1.0, + send_default_pii=True, + ) + + with start_transaction(name="test"): + tool_use = ToolUseBlock(id="tool-1", name="Read", input={"path": "/test.txt"}) + tool_result = ToolResultBlock( + tool_use_id="tool-1", content="file contents here" + ) + span = _create_execute_tool_span(tool_use, tool_result, integration) + span.finish() + + tool_input = span._data[SPANDATA.GEN_AI_TOOL_INPUT] + if isinstance(tool_input, str): + tool_input = json.loads(tool_input) + assert tool_input == {"path": "/test.txt"} + assert span._data[SPANDATA.GEN_AI_TOOL_OUTPUT] == "file contents here" + + +def test_create_execute_tool_span_with_error(sentry_init, integration): + sentry_init( + integrations=[ClaudeAgentSDKIntegration()], + traces_sample_rate=1.0, + send_default_pii=True, + ) + + with start_transaction(name="test"): + tool_use = ToolUseBlock( + id="tool-1", name="Read", input={"path": "/nonexistent.txt"} + ) + tool_result = ToolResultBlock( + tool_use_id="tool-1", + content="Error: file not found", + is_error=True, + ) + span = _create_execute_tool_span(tool_use, tool_result, integration) + span.finish() + + assert span.status == "internal_error" + + +def test_create_execute_tool_span_pii_disabled(sentry_init, integration): + sentry_init( + integrations=[ClaudeAgentSDKIntegration()], + traces_sample_rate=1.0, + send_default_pii=False, + ) + + with start_transaction(name="test"): + tool_use = ToolUseBlock(id="tool-1", name="Read", input={"path": "/test.txt"}) + tool_result = ToolResultBlock(tool_use_id="tool-1", content="file contents") + span = _create_execute_tool_span(tool_use, tool_result, integration) + span.finish() + + assert span._data[SPANDATA.GEN_AI_TOOL_NAME] == "Read" + assert SPANDATA.GEN_AI_TOOL_INPUT not in span._data + assert SPANDATA.GEN_AI_TOOL_OUTPUT not in span._data + + +def test_process_tool_executions_matches_tool_use_and_result(sentry_init, integration): + sentry_init( + integrations=[ClaudeAgentSDKIntegration()], + traces_sample_rate=1.0, + send_default_pii=True, + ) + + with start_transaction(name="test"): + assistant_msg = AssistantMessage( + content=[ + TextBlock(text="Let me read that."), + ToolUseBlock(id="tool-123", name="Read", input={"path": "/test.txt"}), + ToolResultBlock(tool_use_id="tool-123", content="file contents"), + ], + model="test-model", + ) + spans = _process_tool_executions([assistant_msg], integration) + + assert len(spans) == 1 + assert spans[0].description == "execute_tool Read" + + +def test_process_tool_executions_multiple_tools(sentry_init, integration): + sentry_init( + integrations=[ClaudeAgentSDKIntegration()], + traces_sample_rate=1.0, + send_default_pii=True, + ) + + with start_transaction(name="test"): + assistant_msg = AssistantMessage( + content=[ + ToolUseBlock(id="tool-1", name="Read", input={"path": "/a.txt"}), + ToolUseBlock( + id="tool-2", name="Write", input={"path": "/b.txt", "content": "x"} + ), + ToolResultBlock(tool_use_id="tool-1", content="content a"), + ToolResultBlock(tool_use_id="tool-2", content="written"), + ], + model="test-model", + ) + spans = _process_tool_executions([assistant_msg], integration) + + assert len(spans) == 2 + tool_descriptions = {s.description for s in spans} + assert "execute_tool Read" in tool_descriptions + assert "execute_tool Write" in tool_descriptions + + +def test_process_tool_executions_no_tools(sentry_init, integration): + sentry_init( + integrations=[ClaudeAgentSDKIntegration()], + traces_sample_rate=1.0, + send_default_pii=True, + ) + + with start_transaction(name="test"): + assistant_msg = AssistantMessage( + content=[TextBlock(text="Just a text response.")], + model="test-model", + ) + spans = _process_tool_executions([assistant_msg], integration) + + assert len(spans) == 0 + + +@pytest.mark.asyncio +async def test_wrap_query_creates_spans(sentry_init, capture_events): + sentry_init( + integrations=[ClaudeAgentSDKIntegration()], + traces_sample_rate=1.0, + send_default_pii=True, + ) + events = capture_events() + + assistant_msg = AssistantMessage( + content=[TextBlock(text="Hello! I can help you.")], + model="claude-sonnet-4-5-20250929", + ) + result_msg = make_result_message( + usage={"input_tokens": 10, "output_tokens": 20}, + total_cost_usd=0.001, + ) + + async def mock_query(*, prompt, options=None, **kwargs): + yield assistant_msg + yield result_msg + + wrapped = _wrap_query(mock_query) + + with start_transaction(name="claude-agent-sdk-test"): + messages = [] + async for msg in wrapped(prompt="Hello", options=None): + messages.append(msg) + + assert len(messages) == 2 + assert len(events) == 1 + + transaction = events[0] + spans = transaction["spans"] + + invoke_spans = [s for s in spans if s["op"] == OP.GEN_AI_INVOKE_AGENT] + chat_spans = [s for s in spans if s["op"] == OP.GEN_AI_CHAT] + + assert len(invoke_spans) == 1 + assert len(chat_spans) == 1 + + invoke_span = invoke_spans[0] + assert "invoke_agent" in invoke_span["description"] + assert invoke_span["data"][SPANDATA.GEN_AI_OPERATION_NAME] == "invoke_agent" + assert invoke_span["data"][SPANDATA.GEN_AI_SYSTEM] == "claude-agent-sdk-python" + assert invoke_span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 10 + assert invoke_span["data"][SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] == 20 + + chat_span = chat_spans[0] + assert chat_span["data"][SPANDATA.GEN_AI_OPERATION_NAME] == "chat" + assert chat_span["data"][SPANDATA.GEN_AI_SYSTEM] == "claude-agent-sdk-python" + + +@pytest.mark.asyncio +async def test_wrap_query_with_tool_execution(sentry_init, capture_events): + sentry_init( + integrations=[ClaudeAgentSDKIntegration()], + traces_sample_rate=1.0, + send_default_pii=True, + ) + events = capture_events() + + assistant_msg = AssistantMessage( + content=[ + TextBlock(text="Let me read that file."), + ToolUseBlock(id="tool-1", name="Read", input={"path": "/test.txt"}), + ToolResultBlock(tool_use_id="tool-1", content="file contents"), + ], + model="claude-sonnet-4-5-20250929", + ) + result_msg = make_result_message( + usage={"input_tokens": 15, "output_tokens": 25}, + total_cost_usd=0.002, + ) + + async def mock_query(*, prompt, options=None, **kwargs): + yield assistant_msg + yield result_msg + + wrapped = _wrap_query(mock_query) + + with start_transaction(name="claude-agent-sdk-test"): + async for _ in wrapped(prompt="Read /test.txt", options=None): + pass + + assert len(events) == 1 + transaction = events[0] + spans = transaction["spans"] + + tool_spans = [s for s in spans if s["op"] == OP.GEN_AI_EXECUTE_TOOL] + assert len(tool_spans) == 1 + + tool_span = tool_spans[0] + assert tool_span["description"] == "execute_tool Read" + assert tool_span["data"][SPANDATA.GEN_AI_TOOL_NAME] == "Read" + + +@pytest.mark.asyncio +async def test_wrap_query_with_options(sentry_init, capture_events): + sentry_init( + integrations=[ClaudeAgentSDKIntegration()], + traces_sample_rate=1.0, + send_default_pii=True, + ) + events = capture_events() + + assistant_msg = AssistantMessage( + content=[TextBlock(text="I'm helpful!")], + model="claude-opus-4-5-20251101", + ) + result_msg = make_result_message( + usage={"input_tokens": 5, "output_tokens": 10}, + total_cost_usd=0.001, + ) + + async def mock_query(*, prompt, options=None, **kwargs): + yield assistant_msg + yield result_msg + + wrapped = _wrap_query(mock_query) + options = Options( + model="claude-opus-4-5-20251101", + system_prompt="You are helpful.", + ) + + with start_transaction(name="claude-agent-sdk-test"): + async for _ in wrapped(prompt="Hello", options=options): + pass + + transaction = events[0] + spans = transaction["spans"] + + chat_spans = [s for s in spans if s["op"] == OP.GEN_AI_CHAT] + assert len(chat_spans) == 1 + + chat_span = chat_spans[0] + assert ( + chat_span["data"][SPANDATA.GEN_AI_REQUEST_MODEL] == "claude-opus-4-5-20251101" + ) + assert SPANDATA.GEN_AI_REQUEST_MESSAGES in chat_span["data"] + + messages = json.loads(chat_span["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES]) + assert messages[0]["role"] == "system" + assert messages[0]["content"] == "You are helpful." + + +@pytest.mark.asyncio +async def test_wrap_query_pii_disabled(sentry_init, capture_events): + sentry_init( + integrations=[ClaudeAgentSDKIntegration()], + traces_sample_rate=1.0, + send_default_pii=False, + ) + events = capture_events() + + assistant_msg = AssistantMessage( + content=[TextBlock(text="Response text")], + model="claude-sonnet-4-5-20250929", + ) + result_msg = make_result_message( + usage={"input_tokens": 5, "output_tokens": 10}, + total_cost_usd=0.001, + ) + + async def mock_query(*, prompt, options=None, **kwargs): + yield assistant_msg + yield result_msg + + wrapped = _wrap_query(mock_query) + + with start_transaction(name="claude-agent-sdk-test"): + async for _ in wrapped(prompt="Secret prompt", options=None): + pass + + transaction = events[0] + spans = transaction["spans"] + + chat_spans = [s for s in spans if s["op"] == OP.GEN_AI_CHAT] + chat_span = chat_spans[0] + + assert SPANDATA.GEN_AI_REQUEST_MESSAGES not in chat_span["data"] + assert SPANDATA.GEN_AI_RESPONSE_TEXT not in chat_span["data"] + + invoke_spans = [s for s in spans if s["op"] == OP.GEN_AI_INVOKE_AGENT] + invoke_span = invoke_spans[0] + assert invoke_span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 5 + assert invoke_span["data"][SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] == 10 + + +@pytest.mark.asyncio +async def test_wrap_query_handles_exception(sentry_init, capture_events): + sentry_init( + integrations=[ClaudeAgentSDKIntegration()], + traces_sample_rate=1.0, + send_default_pii=True, + ) + events = capture_events() + + async def mock_query(*, prompt, options=None, **kwargs): + yield AssistantMessage(content=[TextBlock(text="Starting...")], model="test") + raise RuntimeError("API error") + + wrapped = _wrap_query(mock_query) + + with pytest.raises(RuntimeError, match="API error"): + with start_transaction(name="claude-agent-sdk-test"): + async for _ in wrapped(prompt="Hello", options=None): + pass + + assert len(events) >= 1