fix(anthropic): Set exception info on span when applicable

Author: alexander-alderman-webb

sentry_sdk/integrations/anthropic.py

@@ -490,16 +490,18 @@ def _wrap_synchronous_message_iterator(
         generator_exit = True
         raise
     finally:
+        exc_info = sys.exc_info()
         with capture_internal_exceptions():
             if not generator_exit and hasattr(stream, "_span"):
-                _finish_streaming_span(
+                _set_streaming_output_data(
                     stream._span,
                     stream._integration,
                     stream._model,
                     stream._usage,
                     stream._content_blocks,
                     stream._response_id,
                 )
+                stream._span.__exit__(*exc_info)
                 del stream._span
 
 
@@ -538,16 +540,18 @@ async def _wrap_asynchronous_message_iterator(
         generator_exit = True
         raise
     finally:
+        exc_info = sys.exc_info()
         with capture_internal_exceptions():
             if not generator_exit and hasattr(stream, "_span"):
-                _finish_streaming_span(
+                _set_streaming_output_data(
                     stream._span,
                     stream._integration,
                     stream._model,
                     stream._usage,
                     stream._content_blocks,
                     stream._response_id,
                 )
+                stream._span.__exit__(*exc_info)
                 del stream._span
 
 
@@ -560,7 +564,6 @@ def _set_output_data(
     cache_read_input_tokens: "int | None",
     cache_write_input_tokens: "int | None",
     content_blocks: "list[Any]",
-    finish_span: bool = False,
     response_id: "str | None" = None,
 ) -> None:
     """
@@ -601,9 +604,6 @@ def _set_output_data(
         input_tokens_cache_write=cache_write_input_tokens,
     )
 
-    if finish_span:
-        span.__exit__(None, None, None)
-
 
 def _sentry_patched_create_common(f: "Any", *args: "Any", **kwargs: "Any") -> "Any":
     integration = kwargs.pop("integration")
@@ -663,9 +663,9 @@ def _sentry_patched_create_common(f: "Any", *args: "Any", **kwargs: "Any") -> "A
                 cache_read_input_tokens=cache_read_input_tokens,
                 cache_write_input_tokens=cache_write_input_tokens,
                 content_blocks=content_blocks,
-                finish_span=True,
                 response_id=getattr(result, "id", None),
             )
+            span.__exit__(None, None, None)
         else:
             span.set_data("unknown_response", True)
             span.__exit__(None, None, None)
@@ -743,7 +743,7 @@ def _accumulate_event_data(
     stream._response_id = response_id
 
 
-def _finish_streaming_span(
+def _set_streaming_output_data(
     span: "Span",
     integration: "AnthropicIntegration",
     model: "Optional[str]",
@@ -752,7 +752,7 @@ def _finish_streaming_span(
     response_id: "Optional[str]",
 ) -> None:
     """
-    Set output attributes on the AI Client Span and end the span.
+    Set output attributes on the AI Client Span.
     """
     # Anthropic's input_tokens excludes cached/cache_write tokens.
     # Normalize to total input tokens for correct cost calculations.
@@ -771,7 +771,6 @@ def _finish_streaming_span(
         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,
         response_id=response_id,
     )
 
@@ -815,14 +814,15 @@ def __next__(self: "Stream") -> "RawMessageStreamEvent":
                 if not hasattr(self, "_span"):
                     raise
 
-                _finish_streaming_span(
+                _set_streaming_output_data(
                     self._span,
                     self._integration,
                     self._model,
                     self._usage,
                     self._content_blocks,
                     self._response_id,
                 )
+                self._span.__exit__(None, None, None)
                 del self._span
             reraise(*exc_info)
 
@@ -847,14 +847,15 @@ def close(self: "Stream") -> None:
             self._span.__exit__(None, None, None)
             return f(self)
 
-        _finish_streaming_span(
+        _set_streaming_output_data(
             self._span,
             self._integration,
             self._model,
             self._usage,
             self._content_blocks,
             self._response_id,
         )
+        self._span.__exit__(None, None, None)
         del self._span
 
         return f(self)
@@ -941,14 +942,15 @@ async def __anext__(self: "AsyncStream") -> "RawMessageStreamEvent":
                 if not hasattr(self, "_span"):
                     raise
 
-                _finish_streaming_span(
+                _set_streaming_output_data(
                     self._span,
                     self._integration,
                     self._model,
                     self._usage,
                     self._content_blocks,
                     self._response_id,
                 )
+                self._span.__exit__(None, None, None)
                 del self._span
             reraise(*exc_info)
 
@@ -973,14 +975,15 @@ async def close(self: "Stream") -> None:
             self._span.__exit__(None, None, None)
             return await f(self)
 
-        _finish_streaming_span(
+        _set_streaming_output_data(
             self._span,
             self._integration,
             self._model,
             self._usage,
             self._content_blocks,
             self._response_id,
         )
+        self._span.__exit__(None, None, None)
         del self._span
 
         return await f(self)
@@ -1104,14 +1107,15 @@ def __next__(self: "MessageStream") -> "MessageStreamEvent":
                 if not hasattr(self, "_span"):
                     raise
 
-                _finish_streaming_span(
+                _set_streaming_output_data(
                     self._span,
                     self._integration,
                     self._model,
                     self._usage,
                     self._content_blocks,
                     self._response_id,
                 )
+                self._span.__exit__(None, None, None)
                 del self._span
             reraise(*exc_info)
 
@@ -1136,14 +1140,15 @@ def close(self: "MessageStream") -> None:
             self._span.__exit__(None, None, None)
             return f(self)
 
-        _finish_streaming_span(
+        _set_streaming_output_data(
             self._span,
             self._integration,
             self._model,
             self._usage,
             self._content_blocks,
             self._response_id,
         )
+        self._span.__exit__(None, None, None)
         del self._span
 
         return f(self)
@@ -1275,14 +1280,15 @@ async def __anext__(self: "AsyncMessageStream") -> "MessageStreamEvent":
                 if not hasattr(self, "_span"):
                     raise
 
-                _finish_streaming_span(
+                _set_streaming_output_data(
                     self._span,
                     self._integration,
                     self._model,
                     self._usage,
                     self._content_blocks,
                     self._response_id,
                 )
+                self._span.__exit__(None, None, None)
                 del self._span
             reraise(*exc_info)
 
@@ -1307,14 +1313,15 @@ async def close(self: "AsyncMessageStream") -> None:
             self._span.__exit__(None, None, None)
             return await f(self)
 
-        _finish_streaming_span(
+        _set_streaming_output_data(
             self._span,
             self._integration,
             self._model,
             self._usage,
             self._content_blocks,
             self._response_id,
         )
+        self._span.__exit__(None, None, None)
         del self._span
 
         return await f(self)