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..106c213352 --- /dev/null +++ b/sentry_sdk/integrations/claude_agent_sdk.py @@ -0,0 +1,422 @@ +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, + AssistantMessage, + ResultMessage, + TextBlock, + ToolUseBlock, + ToolResultBlock, + ) +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" + + +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 isinstance(message, AssistantMessage): + return None + text_parts = [ + block.text for block in getattr(message, "content", []) + if isinstance(block, TextBlock) and hasattr(block, "text") + ] + return "".join(text_parts) if text_parts else None + + +def _extract_tool_calls(message: "Any") -> "Optional[list]": + if not isinstance(message, AssistantMessage): + return None + tool_calls = [] + for block in getattr(message, "content", []): + if isinstance(block, ToolUseBlock): + 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 isinstance(message, AssistantMessage): + 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 isinstance(message, ResultMessage): + data["total_cost"] = getattr(message, "total_cost_usd", None) + usage = getattr(message, "usage", None) + if isinstance(usage, dict): + # Claude Agent SDK returns input_tokens as non-cached tokens only + # For proper cost calculation, we need total input tokens + non_cached_input = usage.get("input_tokens") or 0 + cached_input = usage.get("cache_read_input_tokens") or 0 + # Store total input tokens for cost calculation + data["input_tokens"] = non_cached_input + cached_input + data["output_tokens"] = usage.get("output_tokens") + # Store cached tokens separately for the backend to apply discount pricing + 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"], output_tokens=data["output_tokens"]) + + if data["cached_input_tokens"] is not None: + set_data_normalized(span, SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHED, data["cached_input_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"], 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: "ToolUseBlock", + tool_result: "Optional[ToolResultBlock]", + 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") -> None: + """Create execute_tool spans for tool executions found in messages.""" + tool_uses = {} + tool_results = {} + + for message in messages: + if not isinstance(message, AssistantMessage): + continue + for block in getattr(message, "content", []): + if isinstance(block, ToolUseBlock): + tool_id = getattr(block, "id", None) + if tool_id: + tool_uses[tool_id] = block + elif isinstance(block, ToolResultBlock): + tool_use_id = getattr(block, "tool_use_id", None) + if tool_use_id: + tool_results[tool_use_id] = block + + for tool_id, tool_use in tool_uses.items(): + span = _create_execute_tool_span(tool_use, tool_results.get(tool_id), integration) + span.finish() + + +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..dbe60edb42 --- /dev/null +++ b/tests/integrations/claude_agent_sdk/test_claude_agent_sdk.py @@ -0,0 +1,865 @@ +import pytest +from unittest import mock +from dataclasses import dataclass, field +from typing import Any, 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, + AGENT_NAME, +) + + +# Mock data classes to simulate claude_agent_sdk types +@dataclass +class MockTextBlock: + text: str + type: str = "text" + + +@dataclass +class MockToolUseBlock: + id: str + name: str + input: dict + type: str = "tool_use" + + +@dataclass +class MockToolResultBlock: + tool_use_id: str + content: Optional[str] = None + is_error: bool = False + type: str = "tool_result" + + +@dataclass +class MockAssistantMessage: + content: List[Any] + model: str + error: Optional[str] = None + parent_tool_use_id: Optional[str] = None + + +@dataclass +class MockResultMessage: + subtype: str = "result" + duration_ms: int = 1000 + duration_api_ms: int = 900 + is_error: bool = False + num_turns: int = 1 + session_id: str = "test-session" + total_cost_usd: Optional[float] = 0.005 + usage: Optional[dict] = None + result: Optional[str] = None + structured_output: Any = None + + +@dataclass +class MockClaudeAgentOptions: + model: Optional[str] = None + allowed_tools: Optional[List[str]] = None + system_prompt: Optional[str] = None + max_turns: Optional[int] = None + permission_mode: Optional[str] = None + + +# Fixtures for mock messages +EXAMPLE_ASSISTANT_MESSAGE = MockAssistantMessage( + content=[MockTextBlock(text="Hello! I'm Claude.")], + model="claude-sonnet-4-5-20250929", +) + +EXAMPLE_RESULT_MESSAGE = MockResultMessage( + usage={ + "input_tokens": 10, + "output_tokens": 20, + "cache_read_input_tokens": 100, + }, + total_cost_usd=0.005, +) + +# Module path for patching +INTEGRATION_MODULE = "sentry_sdk.integrations.claude_agent_sdk" + + +def test_extract_text_from_assistant_message(): + """Test extracting text from an AssistantMessage.""" + with mock.patch.multiple( + INTEGRATION_MODULE, + AssistantMessage=MockAssistantMessage, + TextBlock=MockTextBlock, + ): + message = MockAssistantMessage( + content=[MockTextBlock(text="Hello!")], + model="test-model", + ) + text = _extract_text_from_message(message) + assert text == "Hello!" + + +def test_extract_text_from_multiple_blocks(): + """Test extracting text from multiple text blocks.""" + with mock.patch.multiple( + INTEGRATION_MODULE, + AssistantMessage=MockAssistantMessage, + TextBlock=MockTextBlock, + ): + message = MockAssistantMessage( + content=[ + MockTextBlock(text="First. "), + MockTextBlock(text="Second."), + ], + model="test-model", + ) + text = _extract_text_from_message(message) + assert text == "First. Second." + + +def test_extract_tool_calls(): + """Test extracting tool calls from an AssistantMessage.""" + with mock.patch.multiple( + INTEGRATION_MODULE, + AssistantMessage=MockAssistantMessage, + ToolUseBlock=MockToolUseBlock, + ): + message = MockAssistantMessage( + content=[ + MockTextBlock(text="Let me help."), + MockToolUseBlock(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_set_span_input_data_basic(sentry_init): + """Test setting basic input data on a span.""" + 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") + integration = ClaudeAgentSDKIntegration(include_prompts=True) + + _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): + """Test setting input data with options.""" + 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") + integration = ClaudeAgentSDKIntegration(include_prompts=True) + + options = MockClaudeAgentOptions( + 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 + # Check messages include system prompt + 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): + """Test that PII-sensitive data is not captured when PII is disabled.""" + sentry_init( + integrations=[ClaudeAgentSDKIntegration()], + traces_sample_rate=1.0, + send_default_pii=False, # PII disabled + ) + + with start_transaction(name="test") as transaction: + span = transaction.start_child(op="test") + integration = ClaudeAgentSDKIntegration(include_prompts=True) + + _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): + """Test that prompts are not captured when include_prompts is False.""" + 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") + integration = ClaudeAgentSDKIntegration(include_prompts=False) + + _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_output_data_with_messages(sentry_init): + """Test setting output data from messages.""" + sentry_init( + integrations=[ClaudeAgentSDKIntegration()], + traces_sample_rate=1.0, + send_default_pii=True, + ) + + with mock.patch.multiple( + INTEGRATION_MODULE, + AssistantMessage=MockAssistantMessage, + ResultMessage=MockResultMessage, + TextBlock=MockTextBlock, + ): + with start_transaction(name="test") as transaction: + span = transaction.start_child(op="test") + integration = ClaudeAgentSDKIntegration(include_prompts=True) + + messages = [EXAMPLE_ASSISTANT_MESSAGE, EXAMPLE_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): + """Test output data when there's no usage information.""" + sentry_init( + integrations=[ClaudeAgentSDKIntegration()], + traces_sample_rate=1.0, + send_default_pii=True, + ) + + with mock.patch.multiple( + INTEGRATION_MODULE, + AssistantMessage=MockAssistantMessage, + ResultMessage=MockResultMessage, + TextBlock=MockTextBlock, + ): + with start_transaction(name="test") as transaction: + span = transaction.start_child(op="test") + integration = ClaudeAgentSDKIntegration(include_prompts=True) + + result_no_usage = MockResultMessage(usage=None, total_cost_usd=None) + messages = [EXAMPLE_ASSISTANT_MESSAGE, result_no_usage] + _set_span_output_data(span, messages, integration) + + # Should still have model info + assert ( + span._data[SPANDATA.GEN_AI_RESPONSE_MODEL] + == "claude-sonnet-4-5-20250929" + ) + # But no token usage + 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): + """Test output data with tool calls.""" + sentry_init( + integrations=[ClaudeAgentSDKIntegration()], + traces_sample_rate=1.0, + send_default_pii=True, + ) + + with mock.patch.multiple( + INTEGRATION_MODULE, + AssistantMessage=MockAssistantMessage, + ResultMessage=MockResultMessage, + TextBlock=MockTextBlock, + ToolUseBlock=MockToolUseBlock, + ): + with start_transaction(name="test") as transaction: + span = transaction.start_child(op="test") + integration = ClaudeAgentSDKIntegration(include_prompts=True) + + assistant_with_tool = MockAssistantMessage( + content=[ + MockTextBlock(text="Let me read that."), + MockToolUseBlock(id="tool-1", name="Read", input={"path": "/test.txt"}), + ], + model="claude-sonnet-4-5-20250929", + ) + messages = [assistant_with_tool, EXAMPLE_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): + """Test that response text is not captured when PII is disabled.""" + sentry_init( + integrations=[ClaudeAgentSDKIntegration()], + traces_sample_rate=1.0, + send_default_pii=False, # PII disabled + ) + + with mock.patch.multiple( + INTEGRATION_MODULE, + AssistantMessage=MockAssistantMessage, + ResultMessage=MockResultMessage, + TextBlock=MockTextBlock, + ): + with start_transaction(name="test") as transaction: + span = transaction.start_child(op="test") + integration = ClaudeAgentSDKIntegration(include_prompts=True) + + messages = [EXAMPLE_ASSISTANT_MESSAGE, EXAMPLE_RESULT_MESSAGE] + _set_span_output_data(span, messages, integration) + + # Should have model and tokens + assert ( + span._data[SPANDATA.GEN_AI_RESPONSE_MODEL] + == "claude-sonnet-4-5-20250929" + ) + assert span._data[SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 10 + # But not response text + assert SPANDATA.GEN_AI_RESPONSE_TEXT not in span._data + + +def test_integration_identifier(): + """Test that the integration has the correct identifier.""" + integration = ClaudeAgentSDKIntegration() + assert integration.identifier == "claude_agent_sdk" + assert integration.origin == "auto.ai.claude_agent_sdk" + + +def test_integration_include_prompts_default(): + """Test that include_prompts defaults to True.""" + integration = ClaudeAgentSDKIntegration() + assert integration.include_prompts is True + + +def test_integration_include_prompts_false(): + """Test setting include_prompts to 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 +): + """Test the matrix of PII and include_prompts settings.""" + 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): + """Test that request model falls back to response model if not set.""" + sentry_init( + integrations=[ClaudeAgentSDKIntegration()], + traces_sample_rate=1.0, + send_default_pii=True, + ) + + with mock.patch.multiple( + INTEGRATION_MODULE, + AssistantMessage=MockAssistantMessage, + ResultMessage=MockResultMessage, + TextBlock=MockTextBlock, + ): + with start_transaction(name="test") as transaction: + span = transaction.start_child(op="test") + integration = ClaudeAgentSDKIntegration(include_prompts=True) + + # Don't set request model in input + _set_span_input_data(span, "Hello", None, integration) + + # Now set output with response model + messages = [EXAMPLE_ASSISTANT_MESSAGE, EXAMPLE_RESULT_MESSAGE] + _set_span_output_data(span, messages, integration) + + # Request model should be set from response model + 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): + """Test that request model from options is preserved.""" + sentry_init( + integrations=[ClaudeAgentSDKIntegration()], + traces_sample_rate=1.0, + send_default_pii=True, + ) + + with mock.patch.multiple( + INTEGRATION_MODULE, + AssistantMessage=MockAssistantMessage, + ResultMessage=MockResultMessage, + TextBlock=MockTextBlock, + ): + with start_transaction(name="test") as transaction: + span = transaction.start_child(op="test") + integration = ClaudeAgentSDKIntegration(include_prompts=True) + + # Set request model from options + options = MockClaudeAgentOptions(model="claude-opus-4-5-20251101") + _set_span_input_data(span, "Hello", options, integration) + + # Now set output with different response model + messages = [EXAMPLE_ASSISTANT_MESSAGE, EXAMPLE_RESULT_MESSAGE] + _set_span_output_data(span, messages, integration) + + # Request model should be preserved from options + assert ( + span._data[SPANDATA.GEN_AI_REQUEST_MODEL] + == "claude-opus-4-5-20251101" + ) + # Response model should be from response + assert ( + span._data[SPANDATA.GEN_AI_RESPONSE_MODEL] + == "claude-sonnet-4-5-20250929" + ) + + +def test_available_tools_format(sentry_init): + """Test that available tools are formatted correctly.""" + 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") + integration = ClaudeAgentSDKIntegration(include_prompts=True) + + options = MockClaudeAgentOptions(allowed_tools=["Read", "Write", "Bash"]) + _set_span_input_data(span, "Hello", options, integration) + + tools_data = span._data[SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS] + # Should be a JSON string of tool objects + 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): + """Test extraction of cached input tokens.""" + sentry_init( + integrations=[ClaudeAgentSDKIntegration()], + traces_sample_rate=1.0, + send_default_pii=True, + ) + + with mock.patch.multiple( + INTEGRATION_MODULE, + AssistantMessage=MockAssistantMessage, + ResultMessage=MockResultMessage, + TextBlock=MockTextBlock, + ): + with start_transaction(name="test") as transaction: + span = transaction.start_child(op="test") + integration = ClaudeAgentSDKIntegration(include_prompts=True) + + result_with_cache = MockResultMessage( + usage={ + "input_tokens": 5, + "output_tokens": 15, + "cache_read_input_tokens": 500, + }, + total_cost_usd=0.003, + ) + + messages = [EXAMPLE_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_empty_messages_list(sentry_init): + """Test handling of empty messages list.""" + 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") + integration = ClaudeAgentSDKIntegration(include_prompts=True) + + _set_span_output_data(span, [], integration) + + # Should not crash and should not have response data + assert SPANDATA.GEN_AI_RESPONSE_MODEL not in span._data + assert SPANDATA.GEN_AI_RESPONSE_TEXT not in span._data + + +# Tests for invoke_agent spans +def test_start_invoke_agent_span_basic(sentry_init): + """Test starting an invoke_agent span with basic data.""" + sentry_init( + integrations=[ClaudeAgentSDKIntegration()], + traces_sample_rate=1.0, + send_default_pii=True, + ) + + with start_transaction(name="test"): + integration = ClaudeAgentSDKIntegration(include_prompts=True) + span = _start_invoke_agent_span("Hello", None, integration) + + try: + 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 + finally: + span.__exit__(None, None, None) + + +def test_start_invoke_agent_span_with_system_prompt(sentry_init): + """Test invoke_agent span includes system prompt in messages.""" + sentry_init( + integrations=[ClaudeAgentSDKIntegration()], + traces_sample_rate=1.0, + send_default_pii=True, + ) + + with start_transaction(name="test"): + integration = ClaudeAgentSDKIntegration(include_prompts=True) + options = MockClaudeAgentOptions(system_prompt="You are helpful.") + span = _start_invoke_agent_span("Hello", options, integration) + + try: + 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" + finally: + span.__exit__(None, None, None) + + +def test_start_invoke_agent_span_pii_disabled(sentry_init): + """Test invoke_agent span doesn't include messages when PII disabled.""" + sentry_init( + integrations=[ClaudeAgentSDKIntegration()], + traces_sample_rate=1.0, + send_default_pii=False, + ) + + with start_transaction(name="test"): + integration = ClaudeAgentSDKIntegration(include_prompts=True) + span = _start_invoke_agent_span("Hello", None, integration) + + try: + assert span._data[SPANDATA.GEN_AI_SYSTEM] == "claude-agent-sdk-python" + assert SPANDATA.GEN_AI_REQUEST_MESSAGES not in span._data + finally: + span.__exit__(None, None, None) + + +def test_end_invoke_agent_span_aggregates_data(sentry_init): + """Test that end_invoke_agent_span aggregates data from messages.""" + sentry_init( + integrations=[ClaudeAgentSDKIntegration()], + traces_sample_rate=1.0, + send_default_pii=True, + ) + + with mock.patch.multiple( + INTEGRATION_MODULE, + AssistantMessage=MockAssistantMessage, + ResultMessage=MockResultMessage, + TextBlock=MockTextBlock, + ): + with start_transaction(name="test"): + integration = ClaudeAgentSDKIntegration(include_prompts=True) + span = _start_invoke_agent_span("Hello", None, integration) + + messages = [EXAMPLE_ASSISTANT_MESSAGE, EXAMPLE_RESULT_MESSAGE] + _end_invoke_agent_span(span, messages, integration) + + # Check that usage data is set + 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" + + +# Tests for execute_tool spans +def test_create_execute_tool_span_basic(sentry_init): + """Test creating an execute_tool span.""" + sentry_init( + integrations=[ClaudeAgentSDKIntegration()], + traces_sample_rate=1.0, + send_default_pii=True, + ) + + with mock.patch.multiple( + INTEGRATION_MODULE, + ToolUseBlock=MockToolUseBlock, + ToolResultBlock=MockToolResultBlock, + ): + with start_transaction(name="test"): + integration = ClaudeAgentSDKIntegration(include_prompts=True) + tool_use = MockToolUseBlock(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): + """Test execute_tool span includes tool result when available.""" + sentry_init( + integrations=[ClaudeAgentSDKIntegration()], + traces_sample_rate=1.0, + send_default_pii=True, + ) + + with mock.patch.multiple( + INTEGRATION_MODULE, + ToolUseBlock=MockToolUseBlock, + ToolResultBlock=MockToolResultBlock, + ): + with start_transaction(name="test"): + integration = ClaudeAgentSDKIntegration(include_prompts=True) + tool_use = MockToolUseBlock(id="tool-1", name="Read", input={"path": "/test.txt"}) + tool_result = MockToolResultBlock(tool_use_id="tool-1", content="file contents here") + + span = _create_execute_tool_span(tool_use, tool_result, integration) + span.finish() + + # Tool input is stored as JSON string + 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): + """Test execute_tool span sets error status when tool fails.""" + sentry_init( + integrations=[ClaudeAgentSDKIntegration()], + traces_sample_rate=1.0, + send_default_pii=True, + ) + + with mock.patch.multiple( + INTEGRATION_MODULE, + ToolUseBlock=MockToolUseBlock, + ToolResultBlock=MockToolResultBlock, + ): + with start_transaction(name="test"): + integration = ClaudeAgentSDKIntegration(include_prompts=True) + tool_use = MockToolUseBlock(id="tool-1", name="Read", input={"path": "/nonexistent.txt"}) + tool_result = MockToolResultBlock( + 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): + """Test execute_tool span doesn't include input/output when PII disabled.""" + sentry_init( + integrations=[ClaudeAgentSDKIntegration()], + traces_sample_rate=1.0, + send_default_pii=False, + ) + + with mock.patch.multiple( + INTEGRATION_MODULE, + ToolUseBlock=MockToolUseBlock, + ToolResultBlock=MockToolResultBlock, + ): + with start_transaction(name="test"): + integration = ClaudeAgentSDKIntegration(include_prompts=True) + tool_use = MockToolUseBlock(id="tool-1", name="Read", input={"path": "/test.txt"}) + tool_result = MockToolResultBlock(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): + """Test that process_tool_executions matches tool uses with their results.""" + sentry_init( + integrations=[ClaudeAgentSDKIntegration()], + traces_sample_rate=1.0, + send_default_pii=True, + ) + + with mock.patch.multiple( + INTEGRATION_MODULE, + AssistantMessage=MockAssistantMessage, + ToolUseBlock=MockToolUseBlock, + ToolResultBlock=MockToolResultBlock, + ): + with start_transaction(name="test"): + integration = ClaudeAgentSDKIntegration(include_prompts=True) + + # Create messages with tool use and corresponding result + assistant_msg = MockAssistantMessage( + content=[ + MockTextBlock(text="Let me read that."), + MockToolUseBlock(id="tool-123", name="Read", input={"path": "/test.txt"}), + MockToolResultBlock(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): + """Test processing multiple tool executions.""" + sentry_init( + integrations=[ClaudeAgentSDKIntegration()], + traces_sample_rate=1.0, + send_default_pii=True, + ) + + with mock.patch.multiple( + INTEGRATION_MODULE, + AssistantMessage=MockAssistantMessage, + ToolUseBlock=MockToolUseBlock, + ToolResultBlock=MockToolResultBlock, + ): + with start_transaction(name="test"): + integration = ClaudeAgentSDKIntegration(include_prompts=True) + + assistant_msg = MockAssistantMessage( + content=[ + MockToolUseBlock(id="tool-1", name="Read", input={"path": "/a.txt"}), + MockToolUseBlock(id="tool-2", name="Write", input={"path": "/b.txt", "content": "x"}), + MockToolResultBlock(tool_use_id="tool-1", content="content a"), + MockToolResultBlock(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): + """Test that no spans are created when there are no tool uses.""" + sentry_init( + integrations=[ClaudeAgentSDKIntegration()], + traces_sample_rate=1.0, + send_default_pii=True, + ) + + with mock.patch.multiple( + INTEGRATION_MODULE, + AssistantMessage=MockAssistantMessage, + TextBlock=MockTextBlock, + ): + with start_transaction(name="test"): + integration = ClaudeAgentSDKIntegration(include_prompts=True) + + assistant_msg = MockAssistantMessage( + content=[MockTextBlock(text="Just a text response.")], + model="test-model", + ) + + spans = _process_tool_executions([assistant_msg], integration) + + assert len(spans) == 0