fix(anthropic): Close span on GeneratorExit

sentry_sdk/integrations/anthropic.py

@@ -37,8 +37,18 @@
     except ImportError:
         Omit = None
 
+    from anthropic import Stream, AsyncStream
     from anthropic.resources import AsyncMessages, Messages
 
+    from anthropic.types import (
+        MessageStartEvent,
+        MessageDeltaEvent,
+        MessageStopEvent,
+        ContentBlockStartEvent,
+        ContentBlockDeltaEvent,
+        ContentBlockStopEvent,
+    )
+
     if TYPE_CHECKING:
         from anthropic.types import MessageStreamEvent, TextBlockParam
 except ImportError:
@@ -49,6 +59,8 @@
     from sentry_sdk.tracing import Span
     from sentry_sdk._types import TextPart
 
+    from anthropic.types import RawMessageStreamEvent
+
 
 class _RecordedUsage:
     output_tokens: int = 0
@@ -338,6 +350,131 @@
         span.set_data(SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS, safe_serialize(tools))
 
 
+def _wrap_synchronous_message_iterator(
+    iterator: "Iterator[RawMessageStreamEvent]",
+    span: "Span",
+    integration: "AnthropicIntegration",
+) -> "Iterator[RawMessageStreamEvent]":
+    """
+    Sets information received while iterating the response stream on the AI Client Span.
+    Responsible for closing the AI Client Span.
+    """
+
+    model = None
+    usage = _RecordedUsage()
+    content_blocks: "list[str]" = []
+
+    try:
+        for event in iterator:
+            if not isinstance(
+                event,
+                (
+                    MessageStartEvent,
+                    MessageDeltaEvent,
+                    MessageStopEvent,
+                    ContentBlockStartEvent,
+                    ContentBlockDeltaEvent,
+                    ContentBlockStopEvent,
+                ),
+            ):
+                yield event
+                continue
+
+            (
+                model,
+                usage,
+                content_blocks,
+            ) = _collect_ai_data(
+                event,
+                model,
+                usage,
+                content_blocks,
+            )
+            yield event
+    finally:
+        # Anthropic's input_tokens excludes cached/cache_write tokens.
+        # Normalize to total input tokens for correct cost calculations.
+        total_input = (
+            usage.input_tokens
+            + (usage.cache_read_input_tokens or 0)
+            + (usage.cache_write_input_tokens or 0)
+        )
+
+        _set_output_data(
+            span=span,
+            integration=integration,
+            model=model,
+            input_tokens=total_input,
+            output_tokens=usage.output_tokens,
+            cache_read_input_tokens=usage.cache_read_input_tokens,
+            cache_write_input_tokens=usage.cache_write_input_tokens,
+            content_blocks=[{"text": "".join(content_blocks), "type": "text"}],
+            finish_span=True,
+        )
+
+
+async def _wrap_asynchronous_message_iterator(
+    iterator: "AsyncIterator[RawMessageStreamEvent]",
+    span: "Span",
+    integration: "AnthropicIntegration",
+) -> "AsyncIterator[RawMessageStreamEvent]":
+    """
+    Sets information received while iterating the response stream on the AI Client Span.
+    Responsible for closing the AI Client Span.
+    """
+    model = None
+    usage = _RecordedUsage()
+    content_blocks: "list[str]" = []
+
+    try:
+        async for event in iterator:
+            if not isinstance(
+                event,
+                (
+                    MessageStartEvent,
+                    MessageDeltaEvent,
+                    MessageStopEvent,
+                    ContentBlockStartEvent,
+                    ContentBlockDeltaEvent,
+                    ContentBlockStopEvent,
+                ),
+            ):
+                yield event
+                continue
+
+            (
+                model,
+                usage,
+                content_blocks,
+            ) = _collect_ai_data(
+                event,
+                model,
+                usage,
+                content_blocks,
+            )
+            yield event
+    finally:
+        # Anthropic's input_tokens excludes cached/cache_write tokens.
+        # Normalize to total input tokens for correct cost calculations.
+        total_input = (
+            usage.input_tokens
+            + (usage.cache_read_input_tokens or 0)
+            + (usage.cache_write_input_tokens or 0)
+        )
+
+        _set_output_data(
+            span=span,
+            integration=integration,
+            model=model,
+            input_tokens=total_input,
+            output_tokens=usage.output_tokens,
+            cache_read_input_tokens=usage.cache_read_input_tokens,
+            cache_write_input_tokens=usage.cache_write_input_tokens,
+            content_blocks=[{"text": "".join(content_blocks), "type": "text"}],
+            finish_span=True,
+        )
+
+
 def _set_output_data(
     span: "Span",
     integration: "AnthropicIntegration",
@@ -415,6 +552,18 @@
 
     result = yield f, args, kwargs
 
+    if isinstance(result, Stream):
+        result._iterator = _wrap_synchronous_message_iterator(
+            result._iterator, span, integration
+        )
+        return result
+
+    if isinstance(result, AsyncStream):
+        result._iterator = _wrap_asynchronous_message_iterator(
+            result._iterator, span, integration
+        )
+        return result
+
     with capture_internal_exceptions():
         if hasattr(result, "content"):
             (
@@ -444,92 +593,6 @@
                 content_blocks=content_blocks,
                 finish_span=True,
             )
-
-        # Streaming response
-        elif hasattr(result, "_iterator"):
-            old_iterator = result._iterator
-
-            def new_iterator() -> "Iterator[MessageStreamEvent]":
-                model = None
-                usage = _RecordedUsage()
-                content_blocks: "list[str]" = []
-
-                for event in old_iterator:
-                    (
-                        model,
-                        usage,
-                        content_blocks,
-                    ) = _collect_ai_data(
-                        event,
-                        model,
-                        usage,
-                        content_blocks,
-                    )
-                    yield event
-
-                # Anthropic's input_tokens excludes cached/cache_write tokens.
-                # Normalize to total input tokens for correct cost calculations.
-                total_input = (
-                    usage.input_tokens
-                    + (usage.cache_read_input_tokens or 0)
-                    + (usage.cache_write_input_tokens or 0)
-                )
-
-                _set_output_data(
-                    span=span,
-                    integration=integration,
-                    model=model,
-                    input_tokens=total_input,
-                    output_tokens=usage.output_tokens,
-                    cache_read_input_tokens=usage.cache_read_input_tokens,
-                    cache_write_input_tokens=usage.cache_write_input_tokens,
-                    content_blocks=[{"text": "".join(content_blocks), "type": "text"}],
-                    finish_span=True,
-                )
-
-            async def new_iterator_async() -> "AsyncIterator[MessageStreamEvent]":
-                model = None
-                usage = _RecordedUsage()
-                content_blocks: "list[str]" = []
-
-                async for event in old_iterator:
-                    (
-                        model,
-                        usage,
-                        content_blocks,
-                    ) = _collect_ai_data(
-                        event,
-                        model,
-                        usage,
-                        content_blocks,
-                    )
-                    yield event
-
-                # Anthropic's input_tokens excludes cached/cache_write tokens.
-                # Normalize to total input tokens for correct cost calculations.
-                total_input = (
-                    usage.input_tokens
-                    + (usage.cache_read_input_tokens or 0)
-                    + (usage.cache_write_input_tokens or 0)
-                )
-
-                _set_output_data(
-                    span=span,
-                    integration=integration,
-                    model=model,
-                    input_tokens=total_input,
-                    output_tokens=usage.output_tokens,
-                    cache_read_input_tokens=usage.cache_read_input_tokens,
-                    cache_write_input_tokens=usage.cache_write_input_tokens,
-                    content_blocks=[{"text": "".join(content_blocks), "type": "text"}],
-                    finish_span=True,
-                )
-
-            if str(type(result._iterator)) == "<class 'async_generator'>":
-                result._iterator = new_iterator_async()
-            else:
-                result._iterator = new_iterator()
-
         else:
             span.set_data("unknown_response", True)
             span.__exit__(None, None, None)

tests/conftest.py

@@ -57,20 +57,29 @@
     from collections.abc import Iterator
 
 try:
-    from anyio import create_memory_object_stream, create_task_group, EndOfStream
-    from mcp.types import (
-        JSONRPCMessage,
-        JSONRPCNotification,
-        JSONRPCRequest,
-    )
-    from mcp.shared.message import SessionMessage
     from httpx import (
         ASGITransport,
         Request as HttpxRequest,
         Response as HttpxResponse,
         AsyncByteStream,
         AsyncClient,
     )
+except ImportError:
+    ASGITransport = None
+    HttpxRequest = None
+    HttpxResponse = None
+    AsyncByteStream = None
+    AsyncClient = None
+
+
+try:
+    from anyio import create_memory_object_stream, create_task_group, EndOfStream
+    from mcp.types import (
+        JSONRPCMessage,
+        JSONRPCNotification,
+        JSONRPCRequest,
+    )
+    from mcp.shared.message import SessionMessage
 except ImportError:
     create_memory_object_stream = None
     create_task_group = None
@@ -81,12 +90,6 @@
     JSONRPCRequest = None
     SessionMessage = None
 
-    ASGITransport = None
-    HttpxRequest = None
-    HttpxResponse = None
-    AsyncByteStream = None
-    AsyncClient = None
-
 
 SENTRY_EVENT_SCHEMA = "./checkouts/data-schemas/relay/event.schema.json"
 
@@ -1013,6 +1016,39 @@ async def inner(values):
     return inner
 
 
+@pytest.fixture
+def server_side_event_chunks():
+    def inner(events):
+        for event in events:
+            payload = event.model_dump()
+            chunk = f"event: {payload['type']}\ndata: {json.dumps(payload)}\n\n"
+            yield chunk.encode("utf-8")
+
+    return inner
+
+
+@pytest.fixture
+def get_model_response():
+    def inner(response_content, serialize_pydantic=False):
+        model_request = HttpxRequest(
+            "POST",
+            "/responses",
+        )
+
+        if serialize_pydantic:
+            response_content = json.dumps(response_content.model_dump()).encode("utf-8")
+
+        response = HttpxResponse(
+            200,
+            request=model_request,
+            content=response_content,
+        )
+
+        return response
+
+    return inner
+
+
 class MockServerRequestHandler(BaseHTTPRequestHandler):
     def do_GET(self):  # noqa: N802
         # Process an HTTP GET request and return a response.