From 7320fc0287e459aef3e55421a4f9909ba1b8c056 Mon Sep 17 00:00:00 2001 From: Simon Hellmayr Date: Thu, 15 Jan 2026 12:18:30 +0100 Subject: [PATCH 1/8] feat(ai): add claude code agents sdk integration --- sentry_sdk/integrations/claude_code_sdk.py | 391 +++++++++++++++++++++ 1 file changed, 391 insertions(+) create mode 100644 sentry_sdk/integrations/claude_code_sdk.py diff --git a/sentry_sdk/integrations/claude_code_sdk.py b/sentry_sdk/integrations/claude_code_sdk.py new file mode 100644 index 0000000000..d7147a91f3 --- /dev/null +++ b/sentry_sdk/integrations/claude_code_sdk.py @@ -0,0 +1,391 @@ +""" +Sentry integration for Claude Agent SDK (claude-agent-sdk). + +This integration instruments the Claude Agent SDK to capture AI-related +telemetry data, including prompts, responses, token usage, and cost information. + +The integration supports: +- query() function for one-shot queries +- ClaudeSDKClient for interactive sessions + +Usage: + import sentry_sdk + from sentry_sdk.integrations.claude_code_sdk import ClaudeCodeSDKIntegration + + sentry_sdk.init( + dsn="...", + integrations=[ClaudeCodeSDKIntegration()], + traces_sample_rate=1.0, + send_default_pii=True, # Required to capture prompts/responses + ) +""" + +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 +from sentry_sdk.integrations import DidNotEnable, Integration +from sentry_sdk.scope import should_send_default_pii +from sentry_sdk.utils import ( + capture_internal_exceptions, + event_from_exception, +) +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, + ) +except ImportError: + raise DidNotEnable("claude-agent-sdk not installed") + +if TYPE_CHECKING: + from typing import Any, AsyncGenerator, Optional + from sentry_sdk.tracing import Span + + +class ClaudeCodeSDKIntegration(Integration): + """ + Integration for Claude Agent SDK. + + Args: + include_prompts: Whether to include prompts and responses in span data. + Requires send_default_pii=True in Sentry init. Defaults to True. + """ + + identifier = "claude_code_sdk" + origin = f"auto.ai.{identifier}" + + def __init__(self, include_prompts: bool = True) -> None: + self.include_prompts = include_prompts + + @staticmethod + def setup_once() -> None: + # Patch the query function + claude_agent_sdk.query = _wrap_query(original_query) + + # Patch ClaudeSDKClient methods + ClaudeSDKClient._original_query = ClaudeSDKClient.query + ClaudeSDKClient.query = _wrap_client_query(ClaudeSDKClient.query) + + ClaudeSDKClient._original_receive_response = ClaudeSDKClient.receive_response + ClaudeSDKClient.receive_response = _wrap_receive_response( + ClaudeSDKClient.receive_response + ) + + +def _capture_exception(exc: "Any") -> None: + """Capture an exception and set the current span as errored.""" + set_span_errored() + + event, hint = event_from_exception( + exc, + client_options=sentry_sdk.get_client().options, + mechanism={"type": "claude_code_sdk", "handled": False}, + ) + sentry_sdk.capture_event(event, hint=hint) + + +def _set_span_input_data( + span: "Span", + prompt: "str", + options: "Optional[Any]", + integration: "ClaudeCodeSDKIntegration", +) -> None: + """Set input data on the span.""" + set_data_normalized(span, SPANDATA.GEN_AI_SYSTEM, "claude-code") + set_data_normalized(span, SPANDATA.GEN_AI_OPERATION_NAME, "chat") + + # Extract configuration from options if available + if options is not None: + # gen_ai.request.model (required) - will be set from response if not in options + if hasattr(options, "model") and options.model: + set_data_normalized(span, SPANDATA.GEN_AI_REQUEST_MODEL, options.model) + + # gen_ai.request.available_tools (optional) + if hasattr(options, "allowed_tools") and options.allowed_tools: + tools_list = [{"name": tool} for tool in options.allowed_tools] + set_data_normalized( + span, SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS, tools_list, unpack=False + ) + + # gen_ai.request.messages (optional, requires PII) + if hasattr(options, "system_prompt") and options.system_prompt: + if should_send_default_pii() and integration.include_prompts: + set_data_normalized( + span, + SPANDATA.GEN_AI_REQUEST_MESSAGES, + [ + {"role": "system", "content": options.system_prompt}, + {"role": "user", "content": prompt}, + ], + unpack=False, + ) + elif should_send_default_pii() and integration.include_prompts: + set_data_normalized( + span, + SPANDATA.GEN_AI_REQUEST_MESSAGES, + [{"role": "user", "content": prompt}], + unpack=False, + ) + elif should_send_default_pii() and integration.include_prompts: + set_data_normalized( + span, + SPANDATA.GEN_AI_REQUEST_MESSAGES, + [{"role": "user", "content": prompt}], + unpack=False, + ) + + +def _extract_text_from_message(message: "Any") -> "Optional[str]": + """Extract text content from an AssistantMessage.""" + if not isinstance(message, AssistantMessage): + return None + + text_parts = [] + if hasattr(message, "content"): + for block in message.content: + if isinstance(block, TextBlock) and hasattr(block, "text"): + text_parts.append(block.text) + + return "".join(text_parts) if text_parts else None + + +def _extract_tool_calls(message: "Any") -> "Optional[list]": + """Extract tool calls from an AssistantMessage.""" + if not isinstance(message, AssistantMessage): + return None + + tool_calls = [] + if hasattr(message, "content"): + for block in message.content: + if isinstance(block, ToolUseBlock): + tool_call = {"name": getattr(block, "name", "unknown")} + if hasattr(block, "input"): + tool_call["input"] = block.input + tool_calls.append(tool_call) + + return tool_calls if tool_calls else None + + +def _set_span_output_data( + span: "Span", + messages: "list", + integration: "ClaudeCodeSDKIntegration", +) -> None: + """Set output data on the span from collected messages.""" + 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: + response_texts.append(text) + + calls = _extract_tool_calls(message) + if calls: + tool_calls.extend(calls) + + # Extract model from AssistantMessage + if hasattr(message, "model") and message.model and not response_model: + response_model = message.model + + elif isinstance(message, ResultMessage): + if hasattr(message, "total_cost_usd"): + total_cost = message.total_cost_usd + if hasattr(message, "usage") and message.usage: + usage = message.usage + # Usage is a dict with keys like 'input_tokens', 'output_tokens' + if isinstance(usage, dict): + if "input_tokens" in usage: + input_tokens = usage["input_tokens"] + if "output_tokens" in usage: + output_tokens = usage["output_tokens"] + # gen_ai.usage.input_tokens.cached (optional) + if "cache_read_input_tokens" in usage: + cached_input_tokens = usage["cache_read_input_tokens"] + + # gen_ai.response.model (optional, but use to fulfill required gen_ai.request.model) + if response_model: + set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_MODEL, response_model) + # Also set request model if not already set (gen_ai.request.model is required) + # Access span's internal _data dict to check + span_data = getattr(span, "_data", {}) + if SPANDATA.GEN_AI_REQUEST_MODEL not in span_data: + set_data_normalized(span, SPANDATA.GEN_AI_REQUEST_MODEL, response_model) + + # gen_ai.response.text (optional, requires PII) + if response_texts and should_send_default_pii() and integration.include_prompts: + set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, response_texts) + + # gen_ai.response.tool_calls (optional, requires PII) + if tool_calls and should_send_default_pii() and integration.include_prompts: + set_data_normalized( + span, SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS, tool_calls, unpack=False + ) + + # Set token usage if available + # gen_ai.usage.input_tokens, gen_ai.usage.output_tokens, gen_ai.usage.total_tokens (optional) + if input_tokens is not None or output_tokens is not None: + record_token_usage( + span, + input_tokens=input_tokens, + output_tokens=output_tokens, + ) + + # gen_ai.usage.input_tokens.cached (optional) + if cached_input_tokens is not None: + set_data_normalized( + span, SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHED, cached_input_tokens + ) + + # Store cost information in span data + if total_cost is not None: + span.set_data("claude_code.total_cost_usd", total_cost) + + +def _wrap_query(original_func: "Any") -> "Any": + """Wrap the query() async generator function.""" + + @wraps(original_func) + async def wrapper( + *, prompt: str, options: "Optional[Any]" = None, **kwargs: "Any" + ) -> "AsyncGenerator[Any, None]": + integration = sentry_sdk.get_client().get_integration(ClaudeCodeSDKIntegration) + if integration is None: + async for message in original_func(prompt=prompt, options=options, **kwargs): + yield message + return + + model = "" + if options is not None and hasattr(options, "model") and options.model: + model = options.model + + span = get_start_span_function()( + op=OP.GEN_AI_CHAT, + name=f"claude-code query {model}".strip(), + origin=ClaudeCodeSDKIntegration.origin, + ) + span.__enter__() + + with capture_internal_exceptions(): + _set_span_input_data(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(span, collected_messages, integration) + span.__exit__(None, None, None) + + return wrapper + + +def _wrap_client_query(original_method: "Any") -> "Any": + """Wrap the ClaudeSDKClient.query() method.""" + + @wraps(original_method) + async def wrapper(self: "Any", prompt: str, **kwargs: "Any") -> "Any": + integration = sentry_sdk.get_client().get_integration(ClaudeCodeSDKIntegration) + if integration is None: + return await original_method(self, prompt, **kwargs) + + # Store query context on the client for use in receive_response + if not hasattr(self, "_sentry_query_context"): + self._sentry_query_context = {} + + model = "" + if hasattr(self, "_options") and self._options: + if hasattr(self._options, "model") and self._options.model: + model = self._options.model + + span = get_start_span_function()( + op=OP.GEN_AI_CHAT, + name=f"claude-code client {model}".strip(), + origin=ClaudeCodeSDKIntegration.origin, + ) + span.__enter__() + + with capture_internal_exceptions(): + options = getattr(self, "_options", None) + _set_span_input_data(span, prompt, options, integration) + + self._sentry_query_context = { + "span": span, + "integration": integration, + "messages": [], + } + + try: + result = await original_method(self, prompt, **kwargs) + return result + except Exception as exc: + _capture_exception(exc) + # Close span on error + with capture_internal_exceptions(): + _set_span_output_data( + span, self._sentry_query_context.get("messages", []), integration + ) + span.__exit__(None, None, None) + self._sentry_query_context = {} + raise + + return wrapper + + +def _wrap_receive_response(original_method: "Any") -> "Any": + """Wrap the ClaudeSDKClient.receive_response() method.""" + + @wraps(original_method) + async def wrapper(self: "Any", **kwargs: "Any") -> "AsyncGenerator[Any, None]": + integration = sentry_sdk.get_client().get_integration(ClaudeCodeSDKIntegration) + if integration is None: + async for message in original_method(self, **kwargs): + yield message + return + + context = getattr(self, "_sentry_query_context", {}) + span = context.get("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 span is not None: + with capture_internal_exceptions(): + _set_span_output_data(span, messages, stored_integration) + span.__exit__(None, None, None) + self._sentry_query_context = {} + + return wrapper From 52f6f91c813ea440376c2852b2d06f1f68a50f81 Mon Sep 17 00:00:00 2001 From: Simon Hellmayr Date: Thu, 15 Jan 2026 12:25:04 +0100 Subject: [PATCH 2/8] add tests and version --- sentry_sdk/integrations/__init__.py | 2 + sentry_sdk/integrations/claude_code_sdk.py | 6 +- .../integrations/claude_code_sdk/__init__.py | 0 .../claude_code_sdk/test_claude_code_sdk.py | 606 ++++++++++++++++++ 4 files changed, 613 insertions(+), 1 deletion(-) create mode 100644 tests/integrations/claude_code_sdk/__init__.py create mode 100644 tests/integrations/claude_code_sdk/test_claude_code_sdk.py diff --git a/sentry_sdk/integrations/__init__.py b/sentry_sdk/integrations/__init__.py index 9c76dfe471..160e7e7ff3 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_code_sdk.ClaudeCodeSDKIntegration", "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_code_sdk": (0, 1, 0), "cohere": (5, 4, 0), "django": (1, 8), "dramatiq": (1, 9), diff --git a/sentry_sdk/integrations/claude_code_sdk.py b/sentry_sdk/integrations/claude_code_sdk.py index d7147a91f3..7614526dd2 100644 --- a/sentry_sdk/integrations/claude_code_sdk.py +++ b/sentry_sdk/integrations/claude_code_sdk.py @@ -30,11 +30,12 @@ get_start_span_function, ) from sentry_sdk.consts import OP, SPANDATA -from sentry_sdk.integrations import DidNotEnable, Integration +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 @@ -73,6 +74,9 @@ def __init__(self, include_prompts: bool = True) -> None: @staticmethod def setup_once() -> None: + version = package_version("claude_agent_sdk") + _check_minimum_version(ClaudeCodeSDKIntegration, version) + # Patch the query function claude_agent_sdk.query = _wrap_query(original_query) diff --git a/tests/integrations/claude_code_sdk/__init__.py b/tests/integrations/claude_code_sdk/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/integrations/claude_code_sdk/test_claude_code_sdk.py b/tests/integrations/claude_code_sdk/test_claude_code_sdk.py new file mode 100644 index 0000000000..d91ad3a451 --- /dev/null +++ b/tests/integrations/claude_code_sdk/test_claude_code_sdk.py @@ -0,0 +1,606 @@ +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_code_sdk import ( + ClaudeCodeSDKIntegration, + _set_span_input_data, + _set_span_output_data, + _extract_text_from_message, + _extract_tool_calls, +) + + +# Mock data classes to simulate claude_agent_sdk types +@dataclass +class MockTextBlock: + text: str + type: str = "text" + + +@dataclass +class MockToolUseBlock: + name: str + input: dict + type: str = "tool_use" + + +@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, +) + + +def test_extract_text_from_assistant_message(): + """Test extracting text from an AssistantMessage.""" + # Patch the AssistantMessage and TextBlock type checks + with mock.patch( + "sentry_sdk.integrations.claude_code_sdk.AssistantMessage", + MockAssistantMessage, + ): + with mock.patch( + "sentry_sdk.integrations.claude_code_sdk.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( + "sentry_sdk.integrations.claude_code_sdk.AssistantMessage", + MockAssistantMessage, + ): + with mock.patch( + "sentry_sdk.integrations.claude_code_sdk.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( + "sentry_sdk.integrations.claude_code_sdk.AssistantMessage", + MockAssistantMessage, + ): + with mock.patch( + "sentry_sdk.integrations.claude_code_sdk.ToolUseBlock", + MockToolUseBlock, + ): + message = MockAssistantMessage( + content=[ + MockTextBlock(text="Let me help."), + MockToolUseBlock(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=[ClaudeCodeSDKIntegration()], + traces_sample_rate=1.0, + send_default_pii=True, + ) + + with start_transaction(name="test") as transaction: + span = transaction.start_child(op="test") + integration = ClaudeCodeSDKIntegration(include_prompts=True) + + _set_span_input_data(span, "Hello", None, integration) + + assert span._data[SPANDATA.GEN_AI_SYSTEM] == "claude-code" + 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=[ClaudeCodeSDKIntegration()], + traces_sample_rate=1.0, + send_default_pii=True, + ) + + with start_transaction(name="test") as transaction: + span = transaction.start_child(op="test") + integration = ClaudeCodeSDKIntegration(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=[ClaudeCodeSDKIntegration()], + 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 = ClaudeCodeSDKIntegration(include_prompts=True) + + _set_span_input_data(span, "Hello", None, integration) + + assert span._data[SPANDATA.GEN_AI_SYSTEM] == "claude-code" + 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=[ClaudeCodeSDKIntegration(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 = ClaudeCodeSDKIntegration(include_prompts=False) + + _set_span_input_data(span, "Hello", None, integration) + + assert span._data[SPANDATA.GEN_AI_SYSTEM] == "claude-code" + 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=[ClaudeCodeSDKIntegration()], + traces_sample_rate=1.0, + send_default_pii=True, + ) + + with mock.patch( + "sentry_sdk.integrations.claude_code_sdk.AssistantMessage", + MockAssistantMessage, + ): + with mock.patch( + "sentry_sdk.integrations.claude_code_sdk.ResultMessage", + MockResultMessage, + ): + with mock.patch( + "sentry_sdk.integrations.claude_code_sdk.TextBlock", + MockTextBlock, + ): + with start_transaction(name="test") as transaction: + span = transaction.start_child(op="test") + integration = ClaudeCodeSDKIntegration(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=[ClaudeCodeSDKIntegration()], + traces_sample_rate=1.0, + send_default_pii=True, + ) + + with mock.patch( + "sentry_sdk.integrations.claude_code_sdk.AssistantMessage", + MockAssistantMessage, + ): + with mock.patch( + "sentry_sdk.integrations.claude_code_sdk.ResultMessage", + MockResultMessage, + ): + with mock.patch( + "sentry_sdk.integrations.claude_code_sdk.TextBlock", + MockTextBlock, + ): + with start_transaction(name="test") as transaction: + span = transaction.start_child(op="test") + integration = ClaudeCodeSDKIntegration(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=[ClaudeCodeSDKIntegration()], + traces_sample_rate=1.0, + send_default_pii=True, + ) + + with mock.patch( + "sentry_sdk.integrations.claude_code_sdk.AssistantMessage", + MockAssistantMessage, + ): + with mock.patch( + "sentry_sdk.integrations.claude_code_sdk.ResultMessage", + MockResultMessage, + ): + with mock.patch( + "sentry_sdk.integrations.claude_code_sdk.TextBlock", + MockTextBlock, + ): + with mock.patch( + "sentry_sdk.integrations.claude_code_sdk.ToolUseBlock", + MockToolUseBlock, + ): + with start_transaction(name="test") as transaction: + span = transaction.start_child(op="test") + integration = ClaudeCodeSDKIntegration(include_prompts=True) + + assistant_with_tool = MockAssistantMessage( + content=[ + MockTextBlock(text="Let me read that."), + MockToolUseBlock( + 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=[ClaudeCodeSDKIntegration()], + traces_sample_rate=1.0, + send_default_pii=False, # PII disabled + ) + + with mock.patch( + "sentry_sdk.integrations.claude_code_sdk.AssistantMessage", + MockAssistantMessage, + ): + with mock.patch( + "sentry_sdk.integrations.claude_code_sdk.ResultMessage", + MockResultMessage, + ): + with mock.patch( + "sentry_sdk.integrations.claude_code_sdk.TextBlock", + MockTextBlock, + ): + with start_transaction(name="test") as transaction: + span = transaction.start_child(op="test") + integration = ClaudeCodeSDKIntegration(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 = ClaudeCodeSDKIntegration() + assert integration.identifier == "claude_code_sdk" + assert integration.origin == "auto.ai.claude_code_sdk" + + +def test_integration_include_prompts_default(): + """Test that include_prompts defaults to True.""" + integration = ClaudeCodeSDKIntegration() + assert integration.include_prompts is True + + +def test_integration_include_prompts_false(): + """Test setting include_prompts to False.""" + integration = ClaudeCodeSDKIntegration(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=[ClaudeCodeSDKIntegration(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 = ClaudeCodeSDKIntegration(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=[ClaudeCodeSDKIntegration()], + traces_sample_rate=1.0, + send_default_pii=True, + ) + + with mock.patch( + "sentry_sdk.integrations.claude_code_sdk.AssistantMessage", + MockAssistantMessage, + ): + with mock.patch( + "sentry_sdk.integrations.claude_code_sdk.ResultMessage", + MockResultMessage, + ): + with mock.patch( + "sentry_sdk.integrations.claude_code_sdk.TextBlock", + MockTextBlock, + ): + with start_transaction(name="test") as transaction: + span = transaction.start_child(op="test") + integration = ClaudeCodeSDKIntegration(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=[ClaudeCodeSDKIntegration()], + traces_sample_rate=1.0, + send_default_pii=True, + ) + + with mock.patch( + "sentry_sdk.integrations.claude_code_sdk.AssistantMessage", + MockAssistantMessage, + ): + with mock.patch( + "sentry_sdk.integrations.claude_code_sdk.ResultMessage", + MockResultMessage, + ): + with mock.patch( + "sentry_sdk.integrations.claude_code_sdk.TextBlock", + MockTextBlock, + ): + with start_transaction(name="test") as transaction: + span = transaction.start_child(op="test") + integration = ClaudeCodeSDKIntegration(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=[ClaudeCodeSDKIntegration()], + traces_sample_rate=1.0, + send_default_pii=True, + ) + + with start_transaction(name="test") as transaction: + span = transaction.start_child(op="test") + integration = ClaudeCodeSDKIntegration(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=[ClaudeCodeSDKIntegration()], + traces_sample_rate=1.0, + send_default_pii=True, + ) + + with mock.patch( + "sentry_sdk.integrations.claude_code_sdk.AssistantMessage", + MockAssistantMessage, + ): + with mock.patch( + "sentry_sdk.integrations.claude_code_sdk.ResultMessage", + MockResultMessage, + ): + with mock.patch( + "sentry_sdk.integrations.claude_code_sdk.TextBlock", + MockTextBlock, + ): + with start_transaction(name="test") as transaction: + span = transaction.start_child(op="test") + integration = ClaudeCodeSDKIntegration(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=[ClaudeCodeSDKIntegration()], + traces_sample_rate=1.0, + send_default_pii=True, + ) + + with start_transaction(name="test") as transaction: + span = transaction.start_child(op="test") + integration = ClaudeCodeSDKIntegration(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 From 50fa3c448b0b9371ff8a772b225637693ad39c9a Mon Sep 17 00:00:00 2001 From: Simon Hellmayr Date: Thu, 15 Jan 2026 12:30:24 +0100 Subject: [PATCH 3/8] add boilerplate and testing --- pyproject.toml | 4 ++++ scripts/populate_tox/config.py | 7 +++++++ scripts/split_tox_gh_actions/split_tox_gh_actions.py | 1 + setup.py | 1 + 4 files changed, 13 insertions(+) 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..5df094a94a 100644 --- a/scripts/populate_tox/config.py +++ b/scripts/populate_tox/config.py @@ -78,6 +78,13 @@ }, "num_versions": 2, }, + "claude_code_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..a39786c453 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_code_sdk", "openai_agents", "pydantic_ai", ], diff --git a/setup.py b/setup.py index be8e82b26f..e535ba878e 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_code_sdk": ["claude-agent-sdk>=0.1.0"], "clickhouse-driver": ["clickhouse-driver>=0.2.0"], "django": ["django>=1.8"], "falcon": ["falcon>=1.4"], From dbdebfdb9f6f584093b262288f200331aec801be Mon Sep 17 00:00:00 2001 From: Simon Hellmayr Date: Thu, 15 Jan 2026 12:35:03 +0100 Subject: [PATCH 4/8] rename to claude agent sdk --- scripts/populate_tox/config.py | 2 +- .../split_tox_gh_actions.py | 2 +- sentry_sdk/integrations/__init__.py | 4 +- ...claude_code_sdk.py => claude_agent_sdk.py} | 30 ++-- setup.py | 2 +- .../__init__.py | 0 .../test_claude_agent_sdk.py} | 132 +++++++++--------- 7 files changed, 86 insertions(+), 86 deletions(-) rename sentry_sdk/integrations/{claude_code_sdk.py => claude_agent_sdk.py} (94%) rename tests/integrations/{claude_code_sdk => claude_agent_sdk}/__init__.py (100%) rename tests/integrations/{claude_code_sdk/test_claude_code_sdk.py => claude_agent_sdk/test_claude_agent_sdk.py} (80%) diff --git a/scripts/populate_tox/config.py b/scripts/populate_tox/config.py index 5df094a94a..781a590a74 100644 --- a/scripts/populate_tox/config.py +++ b/scripts/populate_tox/config.py @@ -78,7 +78,7 @@ }, "num_versions": 2, }, - "claude_code_sdk": { + "claude_agent_sdk": { "package": "claude-agent-sdk", "deps": { "*": ["pytest-asyncio"], 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 a39786c453..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,7 +74,7 @@ "fastmcp", ], "Agents": [ - "claude_code_sdk", + "claude_agent_sdk", "openai_agents", "pydantic_ai", ], diff --git a/sentry_sdk/integrations/__init__.py b/sentry_sdk/integrations/__init__.py index 160e7e7ff3..2f2520054a 100644 --- a/sentry_sdk/integrations/__init__.py +++ b/sentry_sdk/integrations/__init__.py @@ -67,7 +67,7 @@ def iter_default_integrations( _AUTO_ENABLING_INTEGRATIONS = [ "sentry_sdk.integrations.aiohttp.AioHttpIntegration", "sentry_sdk.integrations.anthropic.AnthropicIntegration", - "sentry_sdk.integrations.claude_code_sdk.ClaudeCodeSDKIntegration", + "sentry_sdk.integrations.claude_agent_sdk.ClaudeCodeSDKIntegration", "sentry_sdk.integrations.ariadne.AriadneIntegration", "sentry_sdk.integrations.arq.ArqIntegration", "sentry_sdk.integrations.asyncpg.AsyncPGIntegration", @@ -128,7 +128,7 @@ def iter_default_integrations( "celery": (4, 4, 7), "chalice": (1, 16, 0), "clickhouse_driver": (0, 2, 0), - "claude_code_sdk": (0, 1, 0), + "claude_agent_sdk": (0, 1, 0), "cohere": (5, 4, 0), "django": (1, 8), "dramatiq": (1, 9), diff --git a/sentry_sdk/integrations/claude_code_sdk.py b/sentry_sdk/integrations/claude_agent_sdk.py similarity index 94% rename from sentry_sdk/integrations/claude_code_sdk.py rename to sentry_sdk/integrations/claude_agent_sdk.py index 7614526dd2..cbdc968b10 100644 --- a/sentry_sdk/integrations/claude_code_sdk.py +++ b/sentry_sdk/integrations/claude_agent_sdk.py @@ -10,7 +10,7 @@ Usage: import sentry_sdk - from sentry_sdk.integrations.claude_code_sdk import ClaudeCodeSDKIntegration + from sentry_sdk.integrations.claude_agent_sdk import ClaudeCodeSDKIntegration sentry_sdk.init( dsn="...", @@ -57,7 +57,7 @@ from sentry_sdk.tracing import Span -class ClaudeCodeSDKIntegration(Integration): +class ClaudeAgentSDKIntegration(Integration): """ Integration for Claude Agent SDK. @@ -66,7 +66,7 @@ class ClaudeCodeSDKIntegration(Integration): Requires send_default_pii=True in Sentry init. Defaults to True. """ - identifier = "claude_code_sdk" + identifier = "claude_agent_sdk" origin = f"auto.ai.{identifier}" def __init__(self, include_prompts: bool = True) -> None: @@ -75,7 +75,7 @@ def __init__(self, include_prompts: bool = True) -> None: @staticmethod def setup_once() -> None: version = package_version("claude_agent_sdk") - _check_minimum_version(ClaudeCodeSDKIntegration, version) + _check_minimum_version(ClaudeAgentSDKIntegration, version) # Patch the query function claude_agent_sdk.query = _wrap_query(original_query) @@ -97,7 +97,7 @@ def _capture_exception(exc: "Any") -> None: event, hint = event_from_exception( exc, client_options=sentry_sdk.get_client().options, - mechanism={"type": "claude_code_sdk", "handled": False}, + mechanism={"type": "claude_agent_sdk", "handled": False}, ) sentry_sdk.capture_event(event, hint=hint) @@ -106,10 +106,10 @@ def _set_span_input_data( span: "Span", prompt: "str", options: "Optional[Any]", - integration: "ClaudeCodeSDKIntegration", + integration: "ClaudeAgentSDKIntegration", ) -> None: """Set input data on the span.""" - set_data_normalized(span, SPANDATA.GEN_AI_SYSTEM, "claude-code") + set_data_normalized(span, SPANDATA.GEN_AI_SYSTEM, "claude-agent-sdk-python") set_data_normalized(span, SPANDATA.GEN_AI_OPERATION_NAME, "chat") # Extract configuration from options if available @@ -187,7 +187,7 @@ def _extract_tool_calls(message: "Any") -> "Optional[list]": def _set_span_output_data( span: "Span", messages: "list", - integration: "ClaudeCodeSDKIntegration", + integration: "ClaudeAgentSDKIntegration", ) -> None: """Set output data on the span from collected messages.""" response_texts = [] @@ -273,7 +273,7 @@ def _wrap_query(original_func: "Any") -> "Any": async def wrapper( *, prompt: str, options: "Optional[Any]" = None, **kwargs: "Any" ) -> "AsyncGenerator[Any, None]": - integration = sentry_sdk.get_client().get_integration(ClaudeCodeSDKIntegration) + 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 @@ -285,8 +285,8 @@ async def wrapper( span = get_start_span_function()( op=OP.GEN_AI_CHAT, - name=f"claude-code query {model}".strip(), - origin=ClaudeCodeSDKIntegration.origin, + name=f"claude-agent-sdk query {model}".strip(), + origin=ClaudeAgentSDKIntegration.origin, ) span.__enter__() @@ -315,7 +315,7 @@ 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(ClaudeCodeSDKIntegration) + integration = sentry_sdk.get_client().get_integration(ClaudeAgentSDKIntegration) if integration is None: return await original_method(self, prompt, **kwargs) @@ -330,8 +330,8 @@ async def wrapper(self: "Any", prompt: str, **kwargs: "Any") -> "Any": span = get_start_span_function()( op=OP.GEN_AI_CHAT, - name=f"claude-code client {model}".strip(), - origin=ClaudeCodeSDKIntegration.origin, + name=f"claude-agent-sdk client {model}".strip(), + origin=ClaudeAgentSDKIntegration.origin, ) span.__enter__() @@ -367,7 +367,7 @@ 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(ClaudeCodeSDKIntegration) + integration = sentry_sdk.get_client().get_integration(ClaudeAgentSDKIntegration) if integration is None: async for message in original_method(self, **kwargs): yield message diff --git a/setup.py b/setup.py index e535ba878e..d36f3dc772 100644 --- a/setup.py +++ b/setup.py @@ -52,7 +52,7 @@ def get_file_text(file_name): "celery": ["celery>=3"], "celery-redbeat": ["celery-redbeat>=2"], "chalice": ["chalice>=1.16.0"], - "claude_code_sdk": ["claude-agent-sdk>=0.1.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_code_sdk/__init__.py b/tests/integrations/claude_agent_sdk/__init__.py similarity index 100% rename from tests/integrations/claude_code_sdk/__init__.py rename to tests/integrations/claude_agent_sdk/__init__.py diff --git a/tests/integrations/claude_code_sdk/test_claude_code_sdk.py b/tests/integrations/claude_agent_sdk/test_claude_agent_sdk.py similarity index 80% rename from tests/integrations/claude_code_sdk/test_claude_code_sdk.py rename to tests/integrations/claude_agent_sdk/test_claude_agent_sdk.py index d91ad3a451..94c6d1098b 100644 --- a/tests/integrations/claude_code_sdk/test_claude_code_sdk.py +++ b/tests/integrations/claude_agent_sdk/test_claude_agent_sdk.py @@ -6,8 +6,8 @@ from sentry_sdk import start_transaction from sentry_sdk.consts import OP, SPANDATA -from sentry_sdk.integrations.claude_code_sdk import ( - ClaudeCodeSDKIntegration, +from sentry_sdk.integrations.claude_agent_sdk import ( + ClaudeAgentSDKIntegration, _set_span_input_data, _set_span_output_data, _extract_text_from_message, @@ -80,11 +80,11 @@ def test_extract_text_from_assistant_message(): """Test extracting text from an AssistantMessage.""" # Patch the AssistantMessage and TextBlock type checks with mock.patch( - "sentry_sdk.integrations.claude_code_sdk.AssistantMessage", + "sentry_sdk.integrations.claude_agent_sdk.AssistantMessage", MockAssistantMessage, ): with mock.patch( - "sentry_sdk.integrations.claude_code_sdk.TextBlock", + "sentry_sdk.integrations.claude_agent_sdk.TextBlock", MockTextBlock, ): message = MockAssistantMessage( @@ -98,11 +98,11 @@ def test_extract_text_from_assistant_message(): def test_extract_text_from_multiple_blocks(): """Test extracting text from multiple text blocks.""" with mock.patch( - "sentry_sdk.integrations.claude_code_sdk.AssistantMessage", + "sentry_sdk.integrations.claude_agent_sdk.AssistantMessage", MockAssistantMessage, ): with mock.patch( - "sentry_sdk.integrations.claude_code_sdk.TextBlock", + "sentry_sdk.integrations.claude_agent_sdk.TextBlock", MockTextBlock, ): message = MockAssistantMessage( @@ -119,11 +119,11 @@ def test_extract_text_from_multiple_blocks(): def test_extract_tool_calls(): """Test extracting tool calls from an AssistantMessage.""" with mock.patch( - "sentry_sdk.integrations.claude_code_sdk.AssistantMessage", + "sentry_sdk.integrations.claude_agent_sdk.AssistantMessage", MockAssistantMessage, ): with mock.patch( - "sentry_sdk.integrations.claude_code_sdk.ToolUseBlock", + "sentry_sdk.integrations.claude_agent_sdk.ToolUseBlock", MockToolUseBlock, ): message = MockAssistantMessage( @@ -142,18 +142,18 @@ def test_extract_tool_calls(): def test_set_span_input_data_basic(sentry_init): """Test setting basic input data on a span.""" sentry_init( - integrations=[ClaudeCodeSDKIntegration()], + 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 = ClaudeCodeSDKIntegration(include_prompts=True) + integration = ClaudeAgentSDKIntegration(include_prompts=True) _set_span_input_data(span, "Hello", None, integration) - assert span._data[SPANDATA.GEN_AI_SYSTEM] == "claude-code" + 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 @@ -161,14 +161,14 @@ def test_set_span_input_data_basic(sentry_init): def test_set_span_input_data_with_options(sentry_init): """Test setting input data with options.""" sentry_init( - integrations=[ClaudeCodeSDKIntegration()], + 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 = ClaudeCodeSDKIntegration(include_prompts=True) + integration = ClaudeAgentSDKIntegration(include_prompts=True) options = MockClaudeAgentOptions( model="claude-opus-4-5-20251101", @@ -191,62 +191,62 @@ def test_set_span_input_data_with_options(sentry_init): 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=[ClaudeCodeSDKIntegration()], + 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 = ClaudeCodeSDKIntegration(include_prompts=True) + integration = ClaudeAgentSDKIntegration(include_prompts=True) _set_span_input_data(span, "Hello", None, integration) - assert span._data[SPANDATA.GEN_AI_SYSTEM] == "claude-code" + 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=[ClaudeCodeSDKIntegration(include_prompts=False)], + 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 = ClaudeCodeSDKIntegration(include_prompts=False) + integration = ClaudeAgentSDKIntegration(include_prompts=False) _set_span_input_data(span, "Hello", None, integration) - assert span._data[SPANDATA.GEN_AI_SYSTEM] == "claude-code" + 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=[ClaudeCodeSDKIntegration()], + integrations=[ClaudeAgentSDKIntegration()], traces_sample_rate=1.0, send_default_pii=True, ) with mock.patch( - "sentry_sdk.integrations.claude_code_sdk.AssistantMessage", + "sentry_sdk.integrations.claude_agent_sdk.AssistantMessage", MockAssistantMessage, ): with mock.patch( - "sentry_sdk.integrations.claude_code_sdk.ResultMessage", + "sentry_sdk.integrations.claude_agent_sdk.ResultMessage", MockResultMessage, ): with mock.patch( - "sentry_sdk.integrations.claude_code_sdk.TextBlock", + "sentry_sdk.integrations.claude_agent_sdk.TextBlock", MockTextBlock, ): with start_transaction(name="test") as transaction: span = transaction.start_child(op="test") - integration = ClaudeCodeSDKIntegration(include_prompts=True) + integration = ClaudeAgentSDKIntegration(include_prompts=True) messages = [EXAMPLE_ASSISTANT_MESSAGE, EXAMPLE_RESULT_MESSAGE] _set_span_output_data(span, messages, integration) @@ -269,26 +269,26 @@ def test_set_span_output_data_with_messages(sentry_init): def test_set_span_output_data_no_usage(sentry_init): """Test output data when there's no usage information.""" sentry_init( - integrations=[ClaudeCodeSDKIntegration()], + integrations=[ClaudeAgentSDKIntegration()], traces_sample_rate=1.0, send_default_pii=True, ) with mock.patch( - "sentry_sdk.integrations.claude_code_sdk.AssistantMessage", + "sentry_sdk.integrations.claude_agent_sdk.AssistantMessage", MockAssistantMessage, ): with mock.patch( - "sentry_sdk.integrations.claude_code_sdk.ResultMessage", + "sentry_sdk.integrations.claude_agent_sdk.ResultMessage", MockResultMessage, ): with mock.patch( - "sentry_sdk.integrations.claude_code_sdk.TextBlock", + "sentry_sdk.integrations.claude_agent_sdk.TextBlock", MockTextBlock, ): with start_transaction(name="test") as transaction: span = transaction.start_child(op="test") - integration = ClaudeCodeSDKIntegration(include_prompts=True) + integration = ClaudeAgentSDKIntegration(include_prompts=True) result_no_usage = MockResultMessage(usage=None, total_cost_usd=None) messages = [EXAMPLE_ASSISTANT_MESSAGE, result_no_usage] @@ -307,30 +307,30 @@ def test_set_span_output_data_no_usage(sentry_init): def test_set_span_output_data_with_tool_calls(sentry_init): """Test output data with tool calls.""" sentry_init( - integrations=[ClaudeCodeSDKIntegration()], + integrations=[ClaudeAgentSDKIntegration()], traces_sample_rate=1.0, send_default_pii=True, ) with mock.patch( - "sentry_sdk.integrations.claude_code_sdk.AssistantMessage", + "sentry_sdk.integrations.claude_agent_sdk.AssistantMessage", MockAssistantMessage, ): with mock.patch( - "sentry_sdk.integrations.claude_code_sdk.ResultMessage", + "sentry_sdk.integrations.claude_agent_sdk.ResultMessage", MockResultMessage, ): with mock.patch( - "sentry_sdk.integrations.claude_code_sdk.TextBlock", + "sentry_sdk.integrations.claude_agent_sdk.TextBlock", MockTextBlock, ): with mock.patch( - "sentry_sdk.integrations.claude_code_sdk.ToolUseBlock", + "sentry_sdk.integrations.claude_agent_sdk.ToolUseBlock", MockToolUseBlock, ): with start_transaction(name="test") as transaction: span = transaction.start_child(op="test") - integration = ClaudeCodeSDKIntegration(include_prompts=True) + integration = ClaudeAgentSDKIntegration(include_prompts=True) assistant_with_tool = MockAssistantMessage( content=[ @@ -350,26 +350,26 @@ def test_set_span_output_data_with_tool_calls(sentry_init): def test_set_span_output_data_pii_disabled(sentry_init): """Test that response text is not captured when PII is disabled.""" sentry_init( - integrations=[ClaudeCodeSDKIntegration()], + integrations=[ClaudeAgentSDKIntegration()], traces_sample_rate=1.0, send_default_pii=False, # PII disabled ) with mock.patch( - "sentry_sdk.integrations.claude_code_sdk.AssistantMessage", + "sentry_sdk.integrations.claude_agent_sdk.AssistantMessage", MockAssistantMessage, ): with mock.patch( - "sentry_sdk.integrations.claude_code_sdk.ResultMessage", + "sentry_sdk.integrations.claude_agent_sdk.ResultMessage", MockResultMessage, ): with mock.patch( - "sentry_sdk.integrations.claude_code_sdk.TextBlock", + "sentry_sdk.integrations.claude_agent_sdk.TextBlock", MockTextBlock, ): with start_transaction(name="test") as transaction: span = transaction.start_child(op="test") - integration = ClaudeCodeSDKIntegration(include_prompts=True) + integration = ClaudeAgentSDKIntegration(include_prompts=True) messages = [EXAMPLE_ASSISTANT_MESSAGE, EXAMPLE_RESULT_MESSAGE] _set_span_output_data(span, messages, integration) @@ -386,20 +386,20 @@ def test_set_span_output_data_pii_disabled(sentry_init): def test_integration_identifier(): """Test that the integration has the correct identifier.""" - integration = ClaudeCodeSDKIntegration() - assert integration.identifier == "claude_code_sdk" - assert integration.origin == "auto.ai.claude_code_sdk" + 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 = ClaudeCodeSDKIntegration() + integration = ClaudeAgentSDKIntegration() assert integration.include_prompts is True def test_integration_include_prompts_false(): """Test setting include_prompts to False.""" - integration = ClaudeCodeSDKIntegration(include_prompts=False) + integration = ClaudeAgentSDKIntegration(include_prompts=False) assert integration.include_prompts is False @@ -417,14 +417,14 @@ def test_pii_and_prompts_matrix( ): """Test the matrix of PII and include_prompts settings.""" sentry_init( - integrations=[ClaudeCodeSDKIntegration(include_prompts=include_prompts)], + 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 = ClaudeCodeSDKIntegration(include_prompts=include_prompts) + integration = ClaudeAgentSDKIntegration(include_prompts=include_prompts) _set_span_input_data(span, "Test prompt", None, integration) @@ -437,26 +437,26 @@ def test_pii_and_prompts_matrix( def test_model_fallback_from_response(sentry_init): """Test that request model falls back to response model if not set.""" sentry_init( - integrations=[ClaudeCodeSDKIntegration()], + integrations=[ClaudeAgentSDKIntegration()], traces_sample_rate=1.0, send_default_pii=True, ) with mock.patch( - "sentry_sdk.integrations.claude_code_sdk.AssistantMessage", + "sentry_sdk.integrations.claude_agent_sdk.AssistantMessage", MockAssistantMessage, ): with mock.patch( - "sentry_sdk.integrations.claude_code_sdk.ResultMessage", + "sentry_sdk.integrations.claude_agent_sdk.ResultMessage", MockResultMessage, ): with mock.patch( - "sentry_sdk.integrations.claude_code_sdk.TextBlock", + "sentry_sdk.integrations.claude_agent_sdk.TextBlock", MockTextBlock, ): with start_transaction(name="test") as transaction: span = transaction.start_child(op="test") - integration = ClaudeCodeSDKIntegration(include_prompts=True) + integration = ClaudeAgentSDKIntegration(include_prompts=True) # Don't set request model in input _set_span_input_data(span, "Hello", None, integration) @@ -479,26 +479,26 @@ def test_model_fallback_from_response(sentry_init): def test_model_from_options_preserved(sentry_init): """Test that request model from options is preserved.""" sentry_init( - integrations=[ClaudeCodeSDKIntegration()], + integrations=[ClaudeAgentSDKIntegration()], traces_sample_rate=1.0, send_default_pii=True, ) with mock.patch( - "sentry_sdk.integrations.claude_code_sdk.AssistantMessage", + "sentry_sdk.integrations.claude_agent_sdk.AssistantMessage", MockAssistantMessage, ): with mock.patch( - "sentry_sdk.integrations.claude_code_sdk.ResultMessage", + "sentry_sdk.integrations.claude_agent_sdk.ResultMessage", MockResultMessage, ): with mock.patch( - "sentry_sdk.integrations.claude_code_sdk.TextBlock", + "sentry_sdk.integrations.claude_agent_sdk.TextBlock", MockTextBlock, ): with start_transaction(name="test") as transaction: span = transaction.start_child(op="test") - integration = ClaudeCodeSDKIntegration(include_prompts=True) + integration = ClaudeAgentSDKIntegration(include_prompts=True) # Set request model from options options = MockClaudeAgentOptions(model="claude-opus-4-5-20251101") @@ -523,14 +523,14 @@ def test_model_from_options_preserved(sentry_init): def test_available_tools_format(sentry_init): """Test that available tools are formatted correctly.""" sentry_init( - integrations=[ClaudeCodeSDKIntegration()], + 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 = ClaudeCodeSDKIntegration(include_prompts=True) + integration = ClaudeAgentSDKIntegration(include_prompts=True) options = MockClaudeAgentOptions(allowed_tools=["Read", "Write", "Bash"]) _set_span_input_data(span, "Hello", options, integration) @@ -548,26 +548,26 @@ def test_available_tools_format(sentry_init): def test_cached_tokens_extraction(sentry_init): """Test extraction of cached input tokens.""" sentry_init( - integrations=[ClaudeCodeSDKIntegration()], + integrations=[ClaudeAgentSDKIntegration()], traces_sample_rate=1.0, send_default_pii=True, ) with mock.patch( - "sentry_sdk.integrations.claude_code_sdk.AssistantMessage", + "sentry_sdk.integrations.claude_agent_sdk.AssistantMessage", MockAssistantMessage, ): with mock.patch( - "sentry_sdk.integrations.claude_code_sdk.ResultMessage", + "sentry_sdk.integrations.claude_agent_sdk.ResultMessage", MockResultMessage, ): with mock.patch( - "sentry_sdk.integrations.claude_code_sdk.TextBlock", + "sentry_sdk.integrations.claude_agent_sdk.TextBlock", MockTextBlock, ): with start_transaction(name="test") as transaction: span = transaction.start_child(op="test") - integration = ClaudeCodeSDKIntegration(include_prompts=True) + integration = ClaudeAgentSDKIntegration(include_prompts=True) result_with_cache = MockResultMessage( usage={ @@ -590,14 +590,14 @@ def test_cached_tokens_extraction(sentry_init): def test_empty_messages_list(sentry_init): """Test handling of empty messages list.""" sentry_init( - integrations=[ClaudeCodeSDKIntegration()], + 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 = ClaudeCodeSDKIntegration(include_prompts=True) + integration = ClaudeAgentSDKIntegration(include_prompts=True) _set_span_output_data(span, [], integration) From 1e2802beacaf69dde9e6e7530ddc3388f126b70b Mon Sep 17 00:00:00 2001 From: Simon Hellmayr Date: Thu, 15 Jan 2026 12:42:17 +0100 Subject: [PATCH 5/8] don't use nested mocks --- sentry_sdk/integrations/__init__.py | 2 +- .../claude_agent_sdk/test_claude_agent_sdk.py | 464 ++++++++---------- 2 files changed, 206 insertions(+), 260 deletions(-) diff --git a/sentry_sdk/integrations/__init__.py b/sentry_sdk/integrations/__init__.py index 2f2520054a..c81b8fa3a8 100644 --- a/sentry_sdk/integrations/__init__.py +++ b/sentry_sdk/integrations/__init__.py @@ -67,7 +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.ClaudeCodeSDKIntegration", + "sentry_sdk.integrations.claude_agent_sdk.ClaudeAgentSDKIntegration", "sentry_sdk.integrations.ariadne.AriadneIntegration", "sentry_sdk.integrations.arq.ArqIntegration", "sentry_sdk.integrations.asyncpg.AsyncPGIntegration", diff --git a/tests/integrations/claude_agent_sdk/test_claude_agent_sdk.py b/tests/integrations/claude_agent_sdk/test_claude_agent_sdk.py index 94c6d1098b..633f15a4d7 100644 --- a/tests/integrations/claude_agent_sdk/test_claude_agent_sdk.py +++ b/tests/integrations/claude_agent_sdk/test_claude_agent_sdk.py @@ -75,68 +75,61 @@ class MockClaudeAgentOptions: 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.""" - # Patch the AssistantMessage and TextBlock type checks - with mock.patch( - "sentry_sdk.integrations.claude_agent_sdk.AssistantMessage", - MockAssistantMessage, + with mock.patch.multiple( + INTEGRATION_MODULE, + AssistantMessage=MockAssistantMessage, + TextBlock=MockTextBlock, ): - with mock.patch( - "sentry_sdk.integrations.claude_agent_sdk.TextBlock", - MockTextBlock, - ): - message = MockAssistantMessage( - content=[MockTextBlock(text="Hello!")], - model="test-model", - ) - text = _extract_text_from_message(message) - assert text == "Hello!" + 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( - "sentry_sdk.integrations.claude_agent_sdk.AssistantMessage", - MockAssistantMessage, + with mock.patch.multiple( + INTEGRATION_MODULE, + AssistantMessage=MockAssistantMessage, + TextBlock=MockTextBlock, ): - with mock.patch( - "sentry_sdk.integrations.claude_agent_sdk.TextBlock", - MockTextBlock, - ): - message = MockAssistantMessage( - content=[ - MockTextBlock(text="First. "), - MockTextBlock(text="Second."), - ], - model="test-model", - ) - text = _extract_text_from_message(message) - assert text == "First. Second." + 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( - "sentry_sdk.integrations.claude_agent_sdk.AssistantMessage", - MockAssistantMessage, + with mock.patch.multiple( + INTEGRATION_MODULE, + AssistantMessage=MockAssistantMessage, + ToolUseBlock=MockToolUseBlock, ): - with mock.patch( - "sentry_sdk.integrations.claude_agent_sdk.ToolUseBlock", - MockToolUseBlock, - ): - message = MockAssistantMessage( - content=[ - MockTextBlock(text="Let me help."), - MockToolUseBlock(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"} + message = MockAssistantMessage( + content=[ + MockTextBlock(text="Let me help."), + MockToolUseBlock(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): @@ -232,38 +225,32 @@ def test_set_span_output_data_with_messages(sentry_init): send_default_pii=True, ) - with mock.patch( - "sentry_sdk.integrations.claude_agent_sdk.AssistantMessage", - MockAssistantMessage, + with mock.patch.multiple( + INTEGRATION_MODULE, + AssistantMessage=MockAssistantMessage, + ResultMessage=MockResultMessage, + TextBlock=MockTextBlock, ): - with mock.patch( - "sentry_sdk.integrations.claude_agent_sdk.ResultMessage", - MockResultMessage, - ): - with mock.patch( - "sentry_sdk.integrations.claude_agent_sdk.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 + 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): @@ -274,34 +261,28 @@ def test_set_span_output_data_no_usage(sentry_init): send_default_pii=True, ) - with mock.patch( - "sentry_sdk.integrations.claude_agent_sdk.AssistantMessage", - MockAssistantMessage, + with mock.patch.multiple( + INTEGRATION_MODULE, + AssistantMessage=MockAssistantMessage, + ResultMessage=MockResultMessage, + TextBlock=MockTextBlock, ): - with mock.patch( - "sentry_sdk.integrations.claude_agent_sdk.ResultMessage", - MockResultMessage, - ): - with mock.patch( - "sentry_sdk.integrations.claude_agent_sdk.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 + 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): @@ -312,39 +293,28 @@ def test_set_span_output_data_with_tool_calls(sentry_init): send_default_pii=True, ) - with mock.patch( - "sentry_sdk.integrations.claude_agent_sdk.AssistantMessage", - MockAssistantMessage, + with mock.patch.multiple( + INTEGRATION_MODULE, + AssistantMessage=MockAssistantMessage, + ResultMessage=MockResultMessage, + TextBlock=MockTextBlock, + ToolUseBlock=MockToolUseBlock, ): - with mock.patch( - "sentry_sdk.integrations.claude_agent_sdk.ResultMessage", - MockResultMessage, - ): - with mock.patch( - "sentry_sdk.integrations.claude_agent_sdk.TextBlock", - MockTextBlock, - ): - with mock.patch( - "sentry_sdk.integrations.claude_agent_sdk.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( - 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 + 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(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): @@ -355,33 +325,27 @@ def test_set_span_output_data_pii_disabled(sentry_init): send_default_pii=False, # PII disabled ) - with mock.patch( - "sentry_sdk.integrations.claude_agent_sdk.AssistantMessage", - MockAssistantMessage, + with mock.patch.multiple( + INTEGRATION_MODULE, + AssistantMessage=MockAssistantMessage, + ResultMessage=MockResultMessage, + TextBlock=MockTextBlock, ): - with mock.patch( - "sentry_sdk.integrations.claude_agent_sdk.ResultMessage", - MockResultMessage, - ): - with mock.patch( - "sentry_sdk.integrations.claude_agent_sdk.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 + 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(): @@ -442,38 +406,32 @@ def test_model_fallback_from_response(sentry_init): send_default_pii=True, ) - with mock.patch( - "sentry_sdk.integrations.claude_agent_sdk.AssistantMessage", - MockAssistantMessage, + with mock.patch.multiple( + INTEGRATION_MODULE, + AssistantMessage=MockAssistantMessage, + ResultMessage=MockResultMessage, + TextBlock=MockTextBlock, ): - with mock.patch( - "sentry_sdk.integrations.claude_agent_sdk.ResultMessage", - MockResultMessage, - ): - with mock.patch( - "sentry_sdk.integrations.claude_agent_sdk.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" - ) + 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): @@ -484,40 +442,34 @@ def test_model_from_options_preserved(sentry_init): send_default_pii=True, ) - with mock.patch( - "sentry_sdk.integrations.claude_agent_sdk.AssistantMessage", - MockAssistantMessage, + with mock.patch.multiple( + INTEGRATION_MODULE, + AssistantMessage=MockAssistantMessage, + ResultMessage=MockResultMessage, + TextBlock=MockTextBlock, ): - with mock.patch( - "sentry_sdk.integrations.claude_agent_sdk.ResultMessage", - MockResultMessage, - ): - with mock.patch( - "sentry_sdk.integrations.claude_agent_sdk.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" - ) + 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): @@ -553,38 +505,32 @@ def test_cached_tokens_extraction(sentry_init): send_default_pii=True, ) - with mock.patch( - "sentry_sdk.integrations.claude_agent_sdk.AssistantMessage", - MockAssistantMessage, + with mock.patch.multiple( + INTEGRATION_MODULE, + AssistantMessage=MockAssistantMessage, + ResultMessage=MockResultMessage, + TextBlock=MockTextBlock, ): - with mock.patch( - "sentry_sdk.integrations.claude_agent_sdk.ResultMessage", - MockResultMessage, - ): - with mock.patch( - "sentry_sdk.integrations.claude_agent_sdk.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 + 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): From bad8e6446a8ed62a8c9489879d9b485abd308429 Mon Sep 17 00:00:00 2001 From: Simon Hellmayr Date: Thu, 15 Jan 2026 13:20:03 +0100 Subject: [PATCH 6/8] add additional span types --- sentry_sdk/integrations/claude_agent_sdk.py | 270 +++++++++++++-- .../claude_agent_sdk/test_claude_agent_sdk.py | 317 +++++++++++++++++- 2 files changed, 555 insertions(+), 32 deletions(-) diff --git a/sentry_sdk/integrations/claude_agent_sdk.py b/sentry_sdk/integrations/claude_agent_sdk.py index cbdc968b10..de0a79b5c6 100644 --- a/sentry_sdk/integrations/claude_agent_sdk.py +++ b/sentry_sdk/integrations/claude_agent_sdk.py @@ -8,13 +8,18 @@ - query() function for one-shot queries - ClaudeSDKClient for interactive sessions +Span hierarchy: +- invoke_agent: Wraps the entire agent invocation (query or client session) + - gen_ai.chat: Individual LLM interactions within the agent invocation + - execute_tool: Tool executions detected in the message stream + Usage: import sentry_sdk - from sentry_sdk.integrations.claude_agent_sdk import ClaudeCodeSDKIntegration + from sentry_sdk.integrations.claude_agent_sdk import ClaudeAgentSDKIntegration sentry_sdk.init( dsn="...", - integrations=[ClaudeCodeSDKIntegration()], + integrations=[ClaudeAgentSDKIntegration()], traces_sample_rate=1.0, send_default_pii=True, # Required to capture prompts/responses ) @@ -29,7 +34,7 @@ set_data_normalized, get_start_span_function, ) -from sentry_sdk.consts import OP, SPANDATA +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 ( @@ -48,14 +53,18 @@ ResultMessage, TextBlock, ToolUseBlock, + ToolResultBlock, ) except ImportError: raise DidNotEnable("claude-agent-sdk not installed") if TYPE_CHECKING: - from typing import Any, AsyncGenerator, Optional + from typing import Any, AsyncGenerator, Optional, Dict, List from sentry_sdk.tracing import Span +# Agent name constant for spans +AGENT_NAME = "claude-agent" + class ClaudeAgentSDKIntegration(Integration): """ @@ -266,8 +275,169 @@ def _set_span_output_data( span.set_data("claude_code.total_cost_usd", total_cost) +def _start_invoke_agent_span( + prompt: "str", + options: "Optional[Any]", + integration: "ClaudeAgentSDKIntegration", +) -> "Span": + """Start an invoke_agent span that wraps the entire agent invocation.""" + 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, "claude-agent-sdk-python") + + # Set request messages if PII enabled + if should_send_default_pii() and integration.include_prompts: + messages = [] + if options is not None and hasattr(options, "system_prompt") and options.system_prompt: + messages.append({"role": "system", "content": options.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: + """End the invoke_agent span with aggregated data from messages.""" + response_texts = [] + total_cost = None + input_tokens = None + output_tokens = None + response_model = None + + for message in messages: + if isinstance(message, AssistantMessage): + text = _extract_text_from_message(message) + if text: + response_texts.append(text) + if hasattr(message, "model") and message.model and not response_model: + response_model = message.model + + elif isinstance(message, ResultMessage): + if hasattr(message, "total_cost_usd"): + total_cost = message.total_cost_usd + if hasattr(message, "usage") and message.usage: + usage = message.usage + if isinstance(usage, dict): + if "input_tokens" in usage: + input_tokens = usage["input_tokens"] + if "output_tokens" in usage: + output_tokens = usage["output_tokens"] + + # Set response text + if response_texts and should_send_default_pii() and integration.include_prompts: + set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, response_texts) + + # Set model info + if response_model: + set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_MODEL, response_model) + set_data_normalized(span, SPANDATA.GEN_AI_REQUEST_MODEL, response_model) + + # Set token usage + if input_tokens is not None or output_tokens is not None: + record_token_usage(span, input_tokens=input_tokens, output_tokens=output_tokens) + + # Set cost + if total_cost is not None: + span.set_data("claude_code.total_cost_usd", total_cost) + + span.__exit__(None, None, None) + + +def _create_execute_tool_span( + tool_use: "ToolUseBlock", + tool_result: "Optional[ToolResultBlock]", + integration: "ClaudeAgentSDKIntegration", +) -> "Span": + """Create an execute_tool span for a tool execution.""" + 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, "claude-agent-sdk-python") + + # Set tool input if PII enabled + if should_send_default_pii() and integration.include_prompts: + tool_input = getattr(tool_use, "input", None) + if tool_input is not None: + set_data_normalized(span, SPANDATA.GEN_AI_TOOL_INPUT, tool_input) + + # Set tool output/result if available + if tool_result is not None: + if should_send_default_pii() and integration.include_prompts: + tool_output = getattr(tool_result, "content", None) + if tool_output is not None: + set_data_normalized(span, SPANDATA.GEN_AI_TOOL_OUTPUT, tool_output) + + # Check for errors + is_error = getattr(tool_result, "is_error", None) + if is_error: + span.set_status(SPANSTATUS.INTERNAL_ERROR) + + return span + + +def _process_tool_executions( + messages: "list", + integration: "ClaudeAgentSDKIntegration", +) -> "List[Span]": + """Process messages to create execute_tool spans for tool executions. + + Tool executions are detected by matching ToolUseBlock with corresponding + ToolResultBlock (matched by tool_use_id). + """ + tool_spans = [] + + # Collect all tool uses and results + tool_uses: "Dict[str, ToolUseBlock]" = {} + tool_results: "Dict[str, ToolResultBlock]" = {} + + for message in messages: + if isinstance(message, AssistantMessage) and hasattr(message, "content"): + for block in 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 + + # Create spans for each tool use + for tool_id, tool_use in tool_uses.items(): + tool_result = tool_results.get(tool_id) + span = _create_execute_tool_span(tool_use, tool_result, integration) + span.finish() + tool_spans.append(span) + + return tool_spans + + def _wrap_query(original_func: "Any") -> "Any": - """Wrap the query() async generator function.""" + """Wrap the query() async generator function. + + Creates an invoke_agent span as the outer span, with a gen_ai.chat span inside. + Tool executions detected in messages will create execute_tool spans. + """ @wraps(original_func) async def wrapper( @@ -283,15 +453,19 @@ async def wrapper( if options is not None and hasattr(options, "model") and options.model: model = options.model - span = get_start_span_function()( + # Start invoke_agent span (outer span) + invoke_span = _start_invoke_agent_span(prompt, options, integration) + + # Start gen_ai.chat span (inner span) + chat_span = get_start_span_function()( op=OP.GEN_AI_CHAT, name=f"claude-agent-sdk query {model}".strip(), origin=ClaudeAgentSDKIntegration.origin, ) - span.__enter__() + chat_span.__enter__() with capture_internal_exceptions(): - _set_span_input_data(span, prompt, options, integration) + _set_span_input_data(chat_span, prompt, options, integration) collected_messages = [] @@ -303,15 +477,28 @@ async def wrapper( _capture_exception(exc) raise finally: + # End chat span with capture_internal_exceptions(): - _set_span_output_data(span, collected_messages, integration) - span.__exit__(None, None, None) + _set_span_output_data(chat_span, collected_messages, integration) + chat_span.__exit__(None, None, None) + + # Create execute_tool spans for any tool executions + with capture_internal_exceptions(): + _process_tool_executions(collected_messages, integration) + + # End invoke_agent span + with capture_internal_exceptions(): + _end_invoke_agent_span(invoke_span, collected_messages, integration) return wrapper def _wrap_client_query(original_method: "Any") -> "Any": - """Wrap the ClaudeSDKClient.query() method.""" + """Wrap the ClaudeSDKClient.query() method. + + Creates an invoke_agent span (outer) and gen_ai.chat span (inner). + The spans are stored on the client instance and completed in receive_response. + """ @wraps(original_method) async def wrapper(self: "Any", prompt: str, **kwargs: "Any") -> "Any": @@ -324,25 +511,31 @@ async def wrapper(self: "Any", prompt: str, **kwargs: "Any") -> "Any": self._sentry_query_context = {} model = "" - if hasattr(self, "_options") and self._options: - if hasattr(self._options, "model") and self._options.model: - model = self._options.model + options = getattr(self, "_options", None) + if options and hasattr(options, "model") and options.model: + model = options.model + + # Start invoke_agent span (outer span) + invoke_span = _start_invoke_agent_span(prompt, options, integration) - span = get_start_span_function()( + # Start gen_ai.chat span (inner span) + chat_span = get_start_span_function()( op=OP.GEN_AI_CHAT, name=f"claude-agent-sdk client {model}".strip(), origin=ClaudeAgentSDKIntegration.origin, ) - span.__enter__() + chat_span.__enter__() with capture_internal_exceptions(): - options = getattr(self, "_options", None) - _set_span_input_data(span, prompt, options, integration) + _set_span_input_data(chat_span, prompt, options, integration) self._sentry_query_context = { - "span": span, + "invoke_span": invoke_span, + "chat_span": chat_span, "integration": integration, "messages": [], + "prompt": prompt, + "options": options, } try: @@ -350,12 +543,13 @@ async def wrapper(self: "Any", prompt: str, **kwargs: "Any") -> "Any": return result except Exception as exc: _capture_exception(exc) - # Close span on error + # Close spans on error + messages = self._sentry_query_context.get("messages", []) with capture_internal_exceptions(): - _set_span_output_data( - span, self._sentry_query_context.get("messages", []), integration - ) - span.__exit__(None, None, None) + _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 @@ -363,7 +557,11 @@ async def wrapper(self: "Any", prompt: str, **kwargs: "Any") -> "Any": def _wrap_receive_response(original_method: "Any") -> "Any": - """Wrap the ClaudeSDKClient.receive_response() method.""" + """Wrap the ClaudeSDKClient.receive_response() method. + + Completes the invoke_agent and chat spans started in client.query(). + Also creates execute_tool spans for any tool executions. + """ @wraps(original_method) async def wrapper(self: "Any", **kwargs: "Any") -> "AsyncGenerator[Any, None]": @@ -374,7 +572,8 @@ async def wrapper(self: "Any", **kwargs: "Any") -> "AsyncGenerator[Any, None]": return context = getattr(self, "_sentry_query_context", {}) - span = context.get("span") + invoke_span = context.get("invoke_span") + chat_span = context.get("chat_span") stored_integration = context.get("integration", integration) messages = context.get("messages", []) @@ -386,10 +585,21 @@ async def wrapper(self: "Any", **kwargs: "Any") -> "AsyncGenerator[Any, None]": _capture_exception(exc) raise finally: - if span is not None: + # End chat span + if chat_span is not None: with capture_internal_exceptions(): - _set_span_output_data(span, messages, stored_integration) - span.__exit__(None, None, None) - self._sentry_query_context = {} + _set_span_output_data(chat_span, messages, stored_integration) + chat_span.__exit__(None, None, None) + + # Create execute_tool spans for any tool executions + with capture_internal_exceptions(): + _process_tool_executions(messages, stored_integration) + + # End invoke_agent span + 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/tests/integrations/claude_agent_sdk/test_claude_agent_sdk.py b/tests/integrations/claude_agent_sdk/test_claude_agent_sdk.py index 633f15a4d7..dbe60edb42 100644 --- a/tests/integrations/claude_agent_sdk/test_claude_agent_sdk.py +++ b/tests/integrations/claude_agent_sdk/test_claude_agent_sdk.py @@ -12,6 +12,11 @@ _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, ) @@ -24,11 +29,20 @@ class MockTextBlock: @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] @@ -122,7 +136,7 @@ def test_extract_tool_calls(): message = MockAssistantMessage( content=[ MockTextBlock(text="Let me help."), - MockToolUseBlock(name="Read", input={"path": "/test.txt"}), + MockToolUseBlock(id="tool-1", name="Read", input={"path": "/test.txt"}), ], model="test-model", ) @@ -307,7 +321,7 @@ def test_set_span_output_data_with_tool_calls(sentry_init): assistant_with_tool = MockAssistantMessage( content=[ MockTextBlock(text="Let me read that."), - MockToolUseBlock(name="Read", input={"path": "/test.txt"}), + MockToolUseBlock(id="tool-1", name="Read", input={"path": "/test.txt"}), ], model="claude-sonnet-4-5-20250929", ) @@ -550,3 +564,302 @@ def test_empty_messages_list(sentry_init): # 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 From 3f869812e47e728e56460cd447f1865990cca632 Mon Sep 17 00:00:00 2001 From: Simon Hellmayr Date: Thu, 15 Jan 2026 13:26:23 +0100 Subject: [PATCH 7/8] remove docstring --- sentry_sdk/integrations/claude_agent_sdk.py | 27 --------------------- 1 file changed, 27 deletions(-) diff --git a/sentry_sdk/integrations/claude_agent_sdk.py b/sentry_sdk/integrations/claude_agent_sdk.py index de0a79b5c6..cb09775147 100644 --- a/sentry_sdk/integrations/claude_agent_sdk.py +++ b/sentry_sdk/integrations/claude_agent_sdk.py @@ -1,30 +1,3 @@ -""" -Sentry integration for Claude Agent SDK (claude-agent-sdk). - -This integration instruments the Claude Agent SDK to capture AI-related -telemetry data, including prompts, responses, token usage, and cost information. - -The integration supports: -- query() function for one-shot queries -- ClaudeSDKClient for interactive sessions - -Span hierarchy: -- invoke_agent: Wraps the entire agent invocation (query or client session) - - gen_ai.chat: Individual LLM interactions within the agent invocation - - execute_tool: Tool executions detected in the message stream - -Usage: - import sentry_sdk - from sentry_sdk.integrations.claude_agent_sdk import ClaudeAgentSDKIntegration - - sentry_sdk.init( - dsn="...", - integrations=[ClaudeAgentSDKIntegration()], - traces_sample_rate=1.0, - send_default_pii=True, # Required to capture prompts/responses - ) -""" - from functools import wraps from typing import TYPE_CHECKING From c8a6111338e5fb600d7142825f335ef8eac37d95 Mon Sep 17 00:00:00 2001 From: Simon Hellmayr Date: Thu, 15 Jan 2026 13:48:57 +0100 Subject: [PATCH 8/8] cleanup --- sentry_sdk/integrations/claude_agent_sdk.py | 434 +++++++------------- 1 file changed, 139 insertions(+), 295 deletions(-) diff --git a/sentry_sdk/integrations/claude_agent_sdk.py b/sentry_sdk/integrations/claude_agent_sdk.py index cb09775147..106c213352 100644 --- a/sentry_sdk/integrations/claude_agent_sdk.py +++ b/sentry_sdk/integrations/claude_agent_sdk.py @@ -3,18 +3,11 @@ 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.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.utils import capture_internal_exceptions, event_from_exception, package_version from sentry_sdk.tracing_utils import set_span_errored try: @@ -32,22 +25,14 @@ raise DidNotEnable("claude-agent-sdk not installed") if TYPE_CHECKING: - from typing import Any, AsyncGenerator, Optional, Dict, List + from typing import Any, AsyncGenerator, Optional from sentry_sdk.tracing import Span -# Agent name constant for spans AGENT_NAME = "claude-agent" +GEN_AI_SYSTEM = "claude-agent-sdk-python" class ClaudeAgentSDKIntegration(Integration): - """ - Integration for Claude Agent SDK. - - Args: - include_prompts: Whether to include prompts and responses in span data. - Requires send_default_pii=True in Sentry init. Defaults to True. - """ - identifier = "claude_agent_sdk" origin = f"auto.ai.{identifier}" @@ -58,24 +43,19 @@ def __init__(self, include_prompts: bool = True) -> None: def setup_once() -> None: version = package_version("claude_agent_sdk") _check_minimum_version(ClaudeAgentSDKIntegration, version) - - # Patch the query function claude_agent_sdk.query = _wrap_query(original_query) - - # Patch ClaudeSDKClient methods - ClaudeSDKClient._original_query = ClaudeSDKClient.query ClaudeSDKClient.query = _wrap_client_query(ClaudeSDKClient.query) - - ClaudeSDKClient._original_receive_response = ClaudeSDKClient.receive_response 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: - """Capture an exception and set the current span as errored.""" set_span_errored() - event, hint = event_from_exception( exc, client_options=sentry_sdk.get_client().options, @@ -86,174 +66,131 @@ def _capture_exception(exc: "Any") -> None: def _set_span_input_data( span: "Span", - prompt: "str", + prompt: str, options: "Optional[Any]", integration: "ClaudeAgentSDKIntegration", ) -> None: - """Set input data on the span.""" - set_data_normalized(span, SPANDATA.GEN_AI_SYSTEM, "claude-agent-sdk-python") + set_data_normalized(span, SPANDATA.GEN_AI_SYSTEM, GEN_AI_SYSTEM) set_data_normalized(span, SPANDATA.GEN_AI_OPERATION_NAME, "chat") - # Extract configuration from options if available if options is not None: - # gen_ai.request.model (required) - will be set from response if not in options - if hasattr(options, "model") and options.model: - set_data_normalized(span, SPANDATA.GEN_AI_REQUEST_MODEL, options.model) - - # gen_ai.request.available_tools (optional) - if hasattr(options, "allowed_tools") and options.allowed_tools: - tools_list = [{"name": tool} for tool in options.allowed_tools] - set_data_normalized( - span, SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS, tools_list, unpack=False - ) - - # gen_ai.request.messages (optional, requires PII) - if hasattr(options, "system_prompt") and options.system_prompt: - if should_send_default_pii() and integration.include_prompts: - set_data_normalized( - span, - SPANDATA.GEN_AI_REQUEST_MESSAGES, - [ - {"role": "system", "content": options.system_prompt}, - {"role": "user", "content": prompt}, - ], - unpack=False, - ) - elif should_send_default_pii() and integration.include_prompts: - set_data_normalized( - span, - SPANDATA.GEN_AI_REQUEST_MESSAGES, - [{"role": "user", "content": prompt}], - unpack=False, - ) - elif should_send_default_pii() and integration.include_prompts: - set_data_normalized( - span, - SPANDATA.GEN_AI_REQUEST_MESSAGES, - [{"role": "user", "content": prompt}], - unpack=False, - ) + 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]": - """Extract text content from an AssistantMessage.""" if not isinstance(message, AssistantMessage): return None - - text_parts = [] - if hasattr(message, "content"): - for block in message.content: - if isinstance(block, TextBlock) and hasattr(block, "text"): - text_parts.append(block.text) - + 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]": - """Extract tool calls from an AssistantMessage.""" if not isinstance(message, AssistantMessage): return None - tool_calls = [] - if hasattr(message, "content"): - for block in message.content: - if isinstance(block, ToolUseBlock): - tool_call = {"name": getattr(block, "name", "unknown")} - if hasattr(block, "input"): - tool_call["input"] = block.input - tool_calls.append(tool_call) - - return tool_calls if tool_calls else None - - -def _set_span_output_data( - span: "Span", - messages: "list", - integration: "ClaudeAgentSDKIntegration", -) -> None: - """Set output data on the span from collected messages.""" - response_texts = [] - tool_calls = [] - total_cost = None - input_tokens = None - output_tokens = None - cached_input_tokens = None - response_model = None + 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: - response_texts.append(text) + data["response_texts"].append(text) calls = _extract_tool_calls(message) if calls: - tool_calls.extend(calls) + data["tool_calls"].extend(calls) - # Extract model from AssistantMessage - if hasattr(message, "model") and message.model and not response_model: - response_model = message.model + if not data["response_model"]: + data["response_model"] = getattr(message, "model", None) elif isinstance(message, ResultMessage): - if hasattr(message, "total_cost_usd"): - total_cost = message.total_cost_usd - if hasattr(message, "usage") and message.usage: - usage = message.usage - # Usage is a dict with keys like 'input_tokens', 'output_tokens' - if isinstance(usage, dict): - if "input_tokens" in usage: - input_tokens = usage["input_tokens"] - if "output_tokens" in usage: - output_tokens = usage["output_tokens"] - # gen_ai.usage.input_tokens.cached (optional) - if "cache_read_input_tokens" in usage: - cached_input_tokens = usage["cache_read_input_tokens"] - - # gen_ai.response.model (optional, but use to fulfill required gen_ai.request.model) - if response_model: - set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_MODEL, response_model) - # Also set request model if not already set (gen_ai.request.model is required) - # Access span's internal _data dict to check - span_data = getattr(span, "_data", {}) - if SPANDATA.GEN_AI_REQUEST_MODEL not in span_data: - set_data_normalized(span, SPANDATA.GEN_AI_REQUEST_MODEL, response_model) - - # gen_ai.response.text (optional, requires PII) - if response_texts and should_send_default_pii() and integration.include_prompts: - set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, response_texts) - - # gen_ai.response.tool_calls (optional, requires PII) - if tool_calls and should_send_default_pii() and integration.include_prompts: - set_data_normalized( - span, SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS, tool_calls, unpack=False - ) + 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 - # Set token usage if available - # gen_ai.usage.input_tokens, gen_ai.usage.output_tokens, gen_ai.usage.total_tokens (optional) - if input_tokens is not None or output_tokens is not None: - record_token_usage( - span, - input_tokens=input_tokens, - output_tokens=output_tokens, - ) - # gen_ai.usage.input_tokens.cached (optional) - if cached_input_tokens is not None: - set_data_normalized( - span, SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHED, cached_input_tokens - ) +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"]) - # Store cost information in span data - if total_cost is not None: - span.set_data("claude_code.total_cost_usd", total_cost) + 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", + prompt: str, options: "Optional[Any]", integration: "ClaudeAgentSDKIntegration", ) -> "Span": - """Start an invoke_agent span that wraps the entire agent invocation.""" span = get_start_span_function()( op=OP.GEN_AI_INVOKE_AGENT, name=f"invoke_agent {AGENT_NAME}", @@ -263,68 +200,38 @@ def _start_invoke_agent_span( 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, "claude-agent-sdk-python") + set_data_normalized(span, SPANDATA.GEN_AI_SYSTEM, GEN_AI_SYSTEM) - # Set request messages if PII enabled - if should_send_default_pii() and integration.include_prompts: + if _should_include_prompts(integration): messages = [] - if options is not None and hasattr(options, "system_prompt") and options.system_prompt: - messages.append({"role": "system", "content": options.system_prompt}) + 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 - ) + set_data_normalized(span, SPANDATA.GEN_AI_REQUEST_MESSAGES, messages, unpack=False) return span def _end_invoke_agent_span( span: "Span", - messages: "list", + messages: list, integration: "ClaudeAgentSDKIntegration", ) -> None: - """End the invoke_agent span with aggregated data from messages.""" - response_texts = [] - total_cost = None - input_tokens = None - output_tokens = None - response_model = None + data = _extract_message_data(messages) - for message in messages: - if isinstance(message, AssistantMessage): - text = _extract_text_from_message(message) - if text: - response_texts.append(text) - if hasattr(message, "model") and message.model and not response_model: - response_model = message.model + if _should_include_prompts(integration) and data["response_texts"]: + set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, data["response_texts"]) - elif isinstance(message, ResultMessage): - if hasattr(message, "total_cost_usd"): - total_cost = message.total_cost_usd - if hasattr(message, "usage") and message.usage: - usage = message.usage - if isinstance(usage, dict): - if "input_tokens" in usage: - input_tokens = usage["input_tokens"] - if "output_tokens" in usage: - output_tokens = usage["output_tokens"] - - # Set response text - if response_texts and should_send_default_pii() and integration.include_prompts: - set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, response_texts) - - # Set model info - if response_model: - set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_MODEL, response_model) - set_data_normalized(span, SPANDATA.GEN_AI_REQUEST_MODEL, response_model) - - # Set token usage - if input_tokens is not None or output_tokens is not None: - record_token_usage(span, input_tokens=input_tokens, output_tokens=output_tokens) - - # Set cost - if total_cost is not None: - span.set_data("claude_code.total_cost_usd", total_cost) + 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) @@ -334,9 +241,7 @@ def _create_execute_tool_span( tool_result: "Optional[ToolResultBlock]", integration: "ClaudeAgentSDKIntegration", ) -> "Span": - """Create an execute_tool span for a tool execution.""" tool_name = getattr(tool_use, "name", "unknown") - span = sentry_sdk.start_span( op=OP.GEN_AI_EXECUTE_TOOL, name=f"execute_tool {tool_name}", @@ -345,73 +250,48 @@ def _create_execute_tool_span( 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, "claude-agent-sdk-python") + set_data_normalized(span, SPANDATA.GEN_AI_SYSTEM, GEN_AI_SYSTEM) - # Set tool input if PII enabled - if should_send_default_pii() and integration.include_prompts: + 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) - # Set tool output/result if available - if tool_result is not None: - if should_send_default_pii() and integration.include_prompts: + 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) - # Check for errors - is_error = getattr(tool_result, "is_error", None) - if is_error: - span.set_status(SPANSTATUS.INTERNAL_ERROR) + 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[Span]": - """Process messages to create execute_tool spans for tool executions. - - Tool executions are detected by matching ToolUseBlock with corresponding - ToolResultBlock (matched by tool_use_id). - """ - tool_spans = [] - - # Collect all tool uses and results - tool_uses: "Dict[str, ToolUseBlock]" = {} - tool_results: "Dict[str, ToolResultBlock]" = {} +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 isinstance(message, AssistantMessage) and hasattr(message, "content"): - for block in 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 - - # Create spans for each tool use + 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(): - tool_result = tool_results.get(tool_id) - span = _create_execute_tool_span(tool_use, tool_result, integration) + span = _create_execute_tool_span(tool_use, tool_results.get(tool_id), integration) span.finish() - tool_spans.append(span) - - return tool_spans def _wrap_query(original_func: "Any") -> "Any": - """Wrap the query() async generator function. - - Creates an invoke_agent span as the outer span, with a gen_ai.chat span inside. - Tool executions detected in messages will create execute_tool spans. - """ - @wraps(original_func) async def wrapper( *, prompt: str, options: "Optional[Any]" = None, **kwargs: "Any" @@ -422,14 +302,9 @@ async def wrapper( yield message return - model = "" - if options is not None and hasattr(options, "model") and options.model: - model = options.model - - # Start invoke_agent span (outer span) + model = getattr(options, "model", "") if options else "" invoke_span = _start_invoke_agent_span(prompt, options, integration) - # Start gen_ai.chat span (inner span) chat_span = get_start_span_function()( op=OP.GEN_AI_CHAT, name=f"claude-agent-sdk query {model}".strip(), @@ -441,7 +316,6 @@ async def wrapper( _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) @@ -450,16 +324,13 @@ async def wrapper( _capture_exception(exc) raise finally: - # End chat span with capture_internal_exceptions(): _set_span_output_data(chat_span, collected_messages, integration) chat_span.__exit__(None, None, None) - # Create execute_tool spans for any tool executions with capture_internal_exceptions(): _process_tool_executions(collected_messages, integration) - # End invoke_agent span with capture_internal_exceptions(): _end_invoke_agent_span(invoke_span, collected_messages, integration) @@ -467,31 +338,17 @@ async def wrapper( def _wrap_client_query(original_method: "Any") -> "Any": - """Wrap the ClaudeSDKClient.query() method. - - Creates an invoke_agent span (outer) and gen_ai.chat span (inner). - The spans are stored on the client instance and completed in receive_response. - """ - @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) - # Store query context on the client for use in receive_response - if not hasattr(self, "_sentry_query_context"): - self._sentry_query_context = {} - - model = "" options = getattr(self, "_options", None) - if options and hasattr(options, "model") and options.model: - model = options.model + model = getattr(options, "model", "") if options else "" - # Start invoke_agent span (outer span) invoke_span = _start_invoke_agent_span(prompt, options, integration) - # Start gen_ai.chat span (inner span) chat_span = get_start_span_function()( op=OP.GEN_AI_CHAT, name=f"claude-agent-sdk client {model}".strip(), @@ -507,16 +364,12 @@ async def wrapper(self: "Any", prompt: str, **kwargs: "Any") -> "Any": "chat_span": chat_span, "integration": integration, "messages": [], - "prompt": prompt, - "options": options, } try: - result = await original_method(self, prompt, **kwargs) - return result + return await original_method(self, prompt, **kwargs) except Exception as exc: _capture_exception(exc) - # Close spans on error messages = self._sentry_query_context.get("messages", []) with capture_internal_exceptions(): _set_span_output_data(chat_span, messages, integration) @@ -530,12 +383,6 @@ async def wrapper(self: "Any", prompt: str, **kwargs: "Any") -> "Any": def _wrap_receive_response(original_method: "Any") -> "Any": - """Wrap the ClaudeSDKClient.receive_response() method. - - Completes the invoke_agent and chat spans started in client.query(). - Also creates execute_tool spans for any tool executions. - """ - @wraps(original_method) async def wrapper(self: "Any", **kwargs: "Any") -> "AsyncGenerator[Any, None]": integration = sentry_sdk.get_client().get_integration(ClaudeAgentSDKIntegration) @@ -558,17 +405,14 @@ async def wrapper(self: "Any", **kwargs: "Any") -> "AsyncGenerator[Any, None]": _capture_exception(exc) raise finally: - # End chat span 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) - # Create execute_tool spans for any tool executions with capture_internal_exceptions(): _process_tool_executions(messages, stored_integration) - # End invoke_agent span if invoke_span is not None: with capture_internal_exceptions(): _end_invoke_agent_span(invoke_span, messages, stored_integration)