Keep track of bedrock stream ending and close spans accordingly

Author: SrdjanLL

instrumentation/opentelemetry-instrumentation-botocore/src/opentelemetry/instrumentation/botocore/extensions/bedrock.py

@@ -606,11 +606,14 @@ def _on_stream_error_callback(
         span: Span,
         exception,
         instrumentor_context: _BotocoreInstrumentorContext,
+        span_ended: bool,
     ):
         span.set_status(Status(StatusCode.ERROR, str(exception)))
         if span.is_recording():
             span.set_attribute(ERROR_TYPE, type(exception).__qualname__)
-        span.end()
+
+        if not span_ended:
+            span.end()
 
         metrics = instrumentor_context.metrics
         metrics_attributes = {
@@ -642,16 +645,17 @@ def on_success(
                 result["stream"], EventStream
             ):
 
-                def stream_done_callback(response):
+                def stream_done_callback(response, span_ended):
                     self._converse_on_success(
                         span, response, instrumentor_context, capture_content
                     )
-                    if span.is_recording():
+
+                    if not span_ended:
                         span.end()
 
-                def stream_error_callback(exception):
+                def stream_error_callback(exception, span_ended):
                     self._on_stream_error_callback(
-                        span, exception, instrumentor_context
+                        span, exception, instrumentor_context, span_ended
                     )
 
                 result["stream"] = ConverseStreamWrapper(
@@ -682,17 +686,17 @@ def stream_error_callback(exception):
         elif self._call_context.operation == "InvokeModelWithResponseStream":
             if "body" in result and isinstance(result["body"], EventStream):
 
-                def invoke_model_stream_done_callback(response):
+                def invoke_model_stream_done_callback(response, span_ended):
                     # the callback gets data formatted as the simpler converse API
                     self._converse_on_success(
                         span, response, instrumentor_context, capture_content
                     )
-                    if span.is_recording():
+                    if not span_ended:
                         span.end()
 
-                def invoke_model_stream_error_callback(exception):
+                def invoke_model_stream_error_callback(exception, span_ended):
                     self._on_stream_error_callback(
-                        span, exception, instrumentor_context
+                        span, exception, instrumentor_context, span_ended
                     )
 
                 result["body"] = InvokeModelWithResponseStreamWrapper(

instrumentation/opentelemetry-instrumentation-botocore/src/opentelemetry/instrumentation/botocore/extensions/bedrock_utils.py

@@ -65,14 +65,15 @@ def __init__(
         self._message = None
         self._content_block = {}
         self._record_message = False
+        self._ended = False
 
     def __iter__(self):
         try:
             for event in self.__wrapped__:
                 self._process_event(event)
                 yield event
         except EventStreamError as exc:
-            self._stream_error_callback(exc)
+            self._handle_stream_error(exc)
             raise
 
     def _process_event(self, event):
@@ -133,15 +134,22 @@ def _process_event(self, event):
 
                 if output_tokens := usage.get("outputTokens"):
                     self._response["usage"]["outputTokens"] = output_tokens
-
-            self._stream_done_callback(self._response)
+            self._complete_stream(self._response)
 
             return
 
     def close(self):
         self.__wrapped__.close()
         # Treat the stream as done to ensure the span end.
-        self._stream_done_callback(self._response)
+        self._complete_stream(self._response)
+
+    def _complete_stream(self, response):
+        self._stream_done_callback(response, self._ended)
+        self._ended = True
+
+    def _handle_stream_error(self, exc):
+        self._stream_error_callback(exc, self._ended)
+        self._ended = True
 
 
 # pylint: disable=abstract-method
@@ -168,19 +176,28 @@ def __init__(
         self._content_block = {}
         self._tool_json_input_buf = ""
         self._record_message = False
+        self._ended = False
 
     def close(self):
         self.__wrapped__.close()
         # Treat the stream as done to ensure the span end.
-        self._stream_done_callback(self._response)
+        self._stream_done_callback(self._response, self._ended)
+
+    def _complete_stream(self, response):
+        self._stream_done_callback(response, self._ended)
+        self._ended = True
+
+    def _handle_stream_error(self, exc):
+        self._stream_error_callback(exc, self._ended)
+        self._ended = True
 
     def __iter__(self):
         try:
             for event in self.__wrapped__:
                 self._process_event(event)
                 yield event
         except EventStreamError as exc:
-            self._stream_error_callback(exc)
+            self._handle_stream_error(exc)
             raise
 
     def _process_event(self, event):
@@ -223,7 +240,7 @@ def _process_amazon_titan_chunk(self, chunk):
             self._response["output"] = {
                 "message": {"content": [{"text": chunk["outputText"]}]}
             }
-            self._stream_done_callback(self._response)
+            self._complete_stream(self._response)
 
     def _process_amazon_nova_chunk(self, chunk):
         # pylint: disable=too-many-branches
@@ -293,7 +310,7 @@ def _process_amazon_nova_chunk(self, chunk):
                 if output_tokens := usage.get("outputTokens"):
                     self._response["usage"]["outputTokens"] = output_tokens
 
-            self._stream_done_callback(self._response)
+            self._complete_stream(self._response)
             return
 
     def _process_anthropic_claude_chunk(self, chunk):
@@ -365,7 +382,7 @@ def _process_anthropic_claude_chunk(self, chunk):
                 self._record_message = False
                 self._message = None
 
-            self._stream_done_callback(self._response)
+            self._complete_stream(self._response)
             return