fixing a bug with tracing paged lists (#41778)

* fixing a bug with tracing paged lists

* updating tests

* updating assests.json

* fixing code checker issues

* updating assest.json

* updating tests

* updating assets.json

* updating version

* updating assets.json

* insignificant change to retrigger build

* updating changelog

Author: M-Hietala

sdk/ai/azure-ai-agents/CHANGELOG.md

@@ -1,5 +1,10 @@
 # Release History
 
+## 1.0.2 (Unreleased)
+
+### Bugs Fixed
+- Fixed a tracing related bug that caused an error when process was ending if messages or run steps were listed and the resulting list was not iterated completely.
+
 ## 1.0.1 (2025-06-09)
 
 ### Bugs Fixed

sdk/ai/azure-ai-agents/assets.json

@@ -2,5 +2,5 @@
   "AssetsRepo": "Azure/azure-sdk-assets",
   "AssetsRepoPrefixPath": "python",
   "TagPrefix": "python/ai/azure-ai-agents",
-  "Tag": "python/ai/azure-ai-agents_dab4ae41b1"
+  "Tag": "python/ai/azure-ai-agents_d7b5157bbd"
 }

sdk/ai/azure-ai-agents/azure/ai/agents/_version.py

@@ -6,4 +6,4 @@
 # Changes may cause incorrect behavior and will be lost if the code is regenerated.
 # --------------------------------------------------------------------------
 
-VERSION = "1.0.1"
+VERSION = "1.0.2"

sdk/ai/azure-ai-agents/azure/ai/agents/telemetry/_ai_agents_instrumentor.py

@@ -317,7 +317,7 @@ def _create_event_attributes(
         if isinstance(end_timestamp, datetime):
             attrs[GEN_AI_RUN_STEP_END_TIMESTAMP] = end_timestamp.isoformat()
         elif end_timestamp:
-            # fallback in case int or string gets passed
+            # fallback in case string or int string gets passed
             attrs[GEN_AI_RUN_STEP_END_TIMESTAMP] = str(end_timestamp)
 
         if run_step_last_error:
@@ -1324,37 +1324,68 @@ def trace_list_messages(self, function, *args, **kwargs):
         server_address = self.get_server_address_from_arg(args[0])
         thread_id = kwargs.get("thread_id")
 
-        span = self.start_list_messages_span(server_address=server_address, thread_id=thread_id)
-
-        return _InstrumentedItemPaged(function(*args, **kwargs), self.add_thread_message_event, span)
+        return _InstrumentedItemPaged(
+            function(*args, **kwargs),
+            start_span_function=self.start_trace_list_messages,
+            item_instrumentation_function=self.add_thread_message_event,
+            server_address=server_address,
+            thread_id=thread_id,
+            run_id=None,
+        )
 
     def trace_list_messages_async(self, function, *args, **kwargs):
         # Note that this method is not async, but it operates on AsyncIterable.
         server_address = self.get_server_address_from_arg(args[0])
         thread_id = kwargs.get("thread_id")
 
-        span = self.start_list_messages_span(server_address=server_address, thread_id=thread_id)
+        return _AsyncInstrumentedItemPaged(
+            function(*args, **kwargs),
+            start_span_function=self.start_trace_list_messages,
+            item_instrumentation_function=self.add_thread_message_event,
+            server_address=server_address,
+            thread_id=thread_id,
+            run_id=None,
+        )
 
-        return _AsyncInstrumentedItemPaged(function(*args, **kwargs), self.add_thread_message_event, span)
+    def start_trace_list_messages(
+        self, server_address: Optional[str] = None, thread_id: Optional[str] = None, run_id: Optional[str] = None
+    ):
+        _ = run_id  # Unused parameter, but kept for compatibility.
+        return self.start_list_messages_span(server_address=server_address, thread_id=thread_id)
 
     def trace_list_run_steps(self, function, *args, **kwargs):
         server_address = self.get_server_address_from_arg(args[0])
         run_id = kwargs.get("run_id")
         thread_id = kwargs.get("thread_id")
 
-        span = self.start_list_run_steps_span(server_address=server_address, run_id=run_id, thread_id=thread_id)
-
-        return _InstrumentedItemPaged(function(*args, **kwargs), self.add_run_step_event, span)
+        return _InstrumentedItemPaged(
+            function(*args, **kwargs),
+            start_span_function=self.start_list_run_steps_span,
+            item_instrumentation_function=self.add_run_step_event,
+            server_address=server_address,
+            thread_id=thread_id,
+            run_id=run_id,
+        )
 
     def trace_list_run_steps_async(self, function, *args, **kwargs):
         # Note that this method is not async, but it operates on AsyncIterable.
         server_address = self.get_server_address_from_arg(args[0])
         run_id = kwargs.get("run_id")
         thread_id = kwargs.get("thread_id")
 
-        span = self.start_list_run_steps_span(server_address=server_address, run_id=run_id, thread_id=thread_id)
+        return _AsyncInstrumentedItemPaged(
+            function(*args, **kwargs),
+            start_span_function=self.start_list_run_steps_span,
+            item_instrumentation_function=self.add_run_step_event,
+            server_address=server_address,
+            thread_id=thread_id,
+            run_id=run_id,
+        )
 
-        return _AsyncInstrumentedItemPaged(function(*args, **kwargs), self.add_run_step_event, span)
+    def start_trace_list_run_steps(
+        self, server_address: Optional[str] = None, thread_id: Optional[str] = None, run_id: Optional[str] = None
+    ):
+        return self.start_list_run_steps_span(server_address=server_address, thread_id=thread_id, run_id=run_id)
 
     def handle_run_stream_exit(self, _function, *args, **kwargs):
         agent_run_stream = args[0]
@@ -2058,6 +2089,7 @@ def on_unhandled_event(self, event_type: str, event_data: Any) -> None:  # type:
         if self.inner_handler:
             return self.inner_handler.on_unhandled_event(event_type, event_data)  # type: ignore
         return super().on_unhandled_event(event_type, event_data)  # type: ignore
+
     # pylint: enable=R1710
 
     def __exit__(self, exc_type, exc_val, exc_tb):

sdk/ai/azure-ai-agents/azure/ai/agents/telemetry/_instrument_paged_wrappers.py

@@ -54,13 +54,19 @@ class _AsyncInstrumentedItemPaged(AsyncItemPaged, _SpanLogger):
     def __init__(
         self,
         async_iter: AsyncItemPaged,
-        instrumentation_fun: Callable[[AbstractSpan, Any], None],
-        span: Optional[AbstractSpan],
+        start_span_function: Callable[[Optional[str], Optional[str], Optional[str]], AbstractSpan],
+        item_instrumentation_function: Callable[[AbstractSpan, Any], None],
+        server_address: Optional[str] = None,
+        thread_id: Optional[str] = None,
+        run_id: Optional[str] = None,
     ) -> None:
         super().__init__()
         self._iter = async_iter
-        self._inst_fun = instrumentation_fun
-        self._span = span
+        self._server_address = server_address
+        self._thread_id = thread_id
+        self._run_id = run_id
+        self._start_span_function = start_span_function
+        self._item_instrumentation_function = item_instrumentation_function
         self._gen: Optional[AsyncIterator[Any]] = None
 
     def __getattr__(self, name: str) -> Any:
@@ -77,20 +83,19 @@ def __getattr__(self, name: str) -> Any:
 
     def __aiter__(self) -> AsyncIterator[Any]:
         async def _gen() -> AsyncIterator[Any]:
-            try:
-                async for val in self._iter:
-                    if self._span is not None:
-                        self.log_to_span_safe(val, self._inst_fun, self._span)
-                    yield val
-            finally:
-                if self._span is not None:
+            async for val in self._iter:
+                span: Optional[AbstractSpan] = self._start_span_function(
+                    self._server_address, self._thread_id, self._run_id
+                )
+                if span is not None:
+                    self.log_to_span_safe(val, self._item_instrumentation_function, span)
                     # We cast None to TracebackType, because traceback is
                     # not used in the downstream code.
-                    self._span.__exit__(None, None, cast(TracebackType, None))
+                    span.__exit__(None, None, cast(TracebackType, None))
+                    span = None
+                yield val
 
         if self._gen is None:
-            if self._span is not None:
-                self._span.__enter__()
             self._gen = _gen()
         return self._gen
 
@@ -101,13 +106,19 @@ class _InstrumentedItemPaged(ItemPaged, _SpanLogger):
     def __init__(
         self,
         iter_val: ItemPaged,
-        instrumentation_fun: Callable[[AbstractSpan, Any], None],
-        span: Optional[AbstractSpan],
+        start_span_function: Callable[[Optional[str], Optional[str], Optional[str]], AbstractSpan],
+        item_instrumentation_function: Callable[[AbstractSpan, Any], None],
+        server_address: Optional[str] = None,
+        thread_id: Optional[str] = None,
+        run_id: Optional[str] = None,
     ) -> None:
         super().__init__()
         self._iter = iter_val
-        self._inst_fun = instrumentation_fun
-        self._span = span
+        self._server_address = server_address
+        self._thread_id = thread_id
+        self._run_id = run_id
+        self._start_span_function = start_span_function
+        self._item_instrumentation_function = item_instrumentation_function
         self._gen: Optional[Iterator[Any]] = None
 
     def __getattr__(self, name: str) -> Any:
@@ -124,19 +135,18 @@ def __getattr__(self, name: str) -> Any:
 
     def __iter__(self) -> Iterator[Any]:
         def _gen() -> Iterator[Any]:
-            try:
-                for val in self._iter:
-                    if self._span is not None:
-                        self.log_to_span_safe(val, self._inst_fun, self._span)
-                    yield val
-            finally:
-                if self._span is not None:
+            for val in self._iter:
+                span: Optional[AbstractSpan] = self._start_span_function(
+                    self._server_address, self._thread_id, self._run_id
+                )
+                if span is not None:
+                    self.log_to_span_safe(val, self._item_instrumentation_function, span)
                     # We cast None to TracebackType, because traceback is
                     # not used in the downstream code.
-                    self._span.__exit__(None, None, cast(TracebackType, None))
+                    span.__exit__(None, None, cast(TracebackType, None))
+                    span = None
+                yield val
 
         if self._gen is None:
-            if self._span is not None:
-                self._span.__enter__()
             self._gen = _gen()
         return self._gen

sdk/ai/azure-ai-agents/tests/conftest.py

@@ -97,6 +97,12 @@ def azure_workspace_triad_sanitizer():
             group_for_replace="1",
         )
 
+        add_general_regex_sanitizer(
+            regex=r"api/projects/([-\w\._\(\)]+)",
+            value=mock_project_scope["project_name"],
+            group_for_replace="1",
+        )
+
     azure_workspace_triad_sanitizer()
 
     add_general_regex_sanitizer(regex=r"/runs/([-\w\._\(\)]+)", value="Sanitized", group_for_replace="1")