fixing tracing for streaming when not using event handler (#41423)

* fixing streaming tracing when event handler is not used

* fixing tracing for streaming when not using event handler

* fixes

* changelog

Author: M-Hietala

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

@@ -4,9 +4,9 @@
 
 ### Bugs Fixed
 
-* `asyncio.gather` is used to make function tool calls in parallel for `async` scenario.
-
-* Adding instrumentation for create_thread_and_run.
+- `asyncio.gather` is used to make function tool calls in parallel for `async` scenario.
+- Adding instrumentation for create_thread_and_run.
+- Fixed a tracing related bug that caused process_thread_run span to not appear when streaming is used without event handler.
 
 ## 1.0.0 (2025-05-15)
 

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

@@ -1379,10 +1379,6 @@ def wrap_handler(
         handler: "Optional[AgentEventHandler]" = None,
         span: "Optional[AbstractSpan]" = None,
     ) -> "Optional[AgentEventHandler]":
-        # Do not create a handler wrapper if we do not have handler in the first place.
-        if not handler:
-            return None
-
         if isinstance(handler, _AgentEventHandlerTraceWrapper):
             return handler
 
@@ -1396,10 +1392,6 @@ def wrap_async_handler(
         handler: "Optional[AsyncAgentEventHandler]" = None,
         span: "Optional[AbstractSpan]" = None,
     ) -> "Optional[AsyncAgentEventHandler]":
-        # Do not create a handler wrapper if we do not have handler in the first place.
-        if not handler:
-            return None
-
         if isinstance(handler, _AsyncAgentEventHandlerTraceWrapper):
             return handler
 
@@ -1975,22 +1967,30 @@ def initialize(
                 response_iterator=response_iterator,
                 submit_tool_outputs=submit_tool_outputs,
             )
+        else:
+            super().initialize(
+                response_iterator=response_iterator,
+                submit_tool_outputs=submit_tool_outputs,
+            )
 
     def __next__(self) -> Any:
         if self.inner_handler:
             event_bytes = self.inner_handler.__next_impl__()
             return self._process_event(event_bytes.decode("utf-8"))
-        return None
+        return super().__next__()
 
     # pylint: disable=R1710
     def on_message_delta(self, delta: "MessageDeltaChunk") -> None:  # type: ignore[func-returns-value]
         if self.inner_handler:
             return self.inner_handler.on_message_delta(delta)  # type: ignore
+        return super().on_message_delta(delta)  # type: ignore
 
     def on_thread_message(self, message: "ThreadMessage") -> None:  # type: ignore[func-returns-value]
         retval = None
         if self.inner_handler:
             retval = self.inner_handler.on_thread_message(message)  # type: ignore
+        else:
+            retval = super().on_thread_message(message)  # pylint: disable=assignment-from-none # type: ignore
 
         if message.status in {"completed", "incomplete"}:
             self.last_message = message
@@ -2002,6 +2002,8 @@ def on_thread_run(self, run: "ThreadRun") -> None:  # type: ignore[func-returns-
 
         if self.inner_handler:
             retval = self.inner_handler.on_thread_run(run)  # type: ignore
+        else:
+            retval = super().on_thread_run(run)  # pylint: disable=assignment-from-none # type: ignore
         self.last_run = run
 
         return retval  # type: ignore
@@ -2010,6 +2012,8 @@ def on_run_step(self, step: "RunStep") -> None:  # type: ignore[func-returns-val
         retval = None
         if self.inner_handler:
             retval = self.inner_handler.on_run_step(step)  # type: ignore
+        else:
+            retval = super().on_run_step(step)  # pylint: disable=assignment-from-none # type: ignore
 
         if (
             step.type == "tool_calls"
@@ -2035,20 +2039,23 @@ def on_run_step(self, step: "RunStep") -> None:  # type: ignore[func-returns-val
     def on_run_step_delta(self, delta: "RunStepDeltaChunk") -> None:  # type: ignore[func-returns-value]
         if self.inner_handler:
             return self.inner_handler.on_run_step_delta(delta)  # type: ignore
+        return super().on_run_step_delta(delta)  # type: ignore
 
     def on_error(self, data: str) -> None:  # type: ignore[func-returns-value]
         if self.inner_handler:
             return self.inner_handler.on_error(data)  # type: ignore
+        return super().on_error(data)  # type: ignore
 
     def on_done(self) -> None:  # type: ignore[func-returns-value]
         if self.inner_handler:
             return self.inner_handler.on_done()  # type: ignore
+        return super().on_done()  # type: ignore
         # it could be called multiple tines (for each step) __exit__
 
     def on_unhandled_event(self, event_type: str, event_data: Any) -> None:  # type: ignore[func-returns-value]
         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):
@@ -2093,23 +2100,33 @@ def initialize(
                 response_iterator=response_iterator,
                 submit_tool_outputs=submit_tool_outputs,
             )
+        else:
+            super().initialize(
+                response_iterator=response_iterator,
+                submit_tool_outputs=submit_tool_outputs,
+            )
 
     # cspell:disable-next-line
     async def __anext__(self) -> Any:
         if self.inner_handler:
             # cspell:disable-next-line
             event_bytes = await self.inner_handler.__anext_impl__()
             return await self._process_event(event_bytes.decode("utf-8"))
+        # cspell:disable-next-line
+        return await super().__anext__()
 
     # pylint: disable=R1710
     async def on_message_delta(self, delta: "MessageDeltaChunk") -> None:  # type: ignore[func-returns-value]
         if self.inner_handler:
             return await self.inner_handler.on_message_delta(delta)  # type: ignore
+        return await super().on_message_delta(delta)  # type: ignore
 
     async def on_thread_message(self, message: "ThreadMessage") -> None:  # type: ignore[func-returns-value]
         retval = None
         if self.inner_handler:
             retval = await self.inner_handler.on_thread_message(message)  # type: ignore
+        else:
+            retval = await super().on_thread_message(message)  # type: ignore
 
         if message.status in {"completed", "incomplete"}:
             self.last_message = message
@@ -2121,6 +2138,8 @@ async def on_thread_run(self, run: "ThreadRun") -> None:  # type: ignore[func-re
 
         if self.inner_handler:
             retval = await self.inner_handler.on_thread_run(run)  # type: ignore
+        else:
+            retval = await super().on_thread_run(run)  # type: ignore
         self.last_run = run
 
         return retval  # type: ignore
@@ -2129,6 +2148,8 @@ async def on_run_step(self, step: "RunStep") -> None:  # type: ignore[func-retur
         retval = None
         if self.inner_handler:
             retval = await self.inner_handler.on_run_step(step)  # type: ignore
+        else:
+            retval = await super().on_run_step(step)  # type: ignore
 
         if (
             step.type == "tool_calls"
@@ -2154,19 +2175,23 @@ async def on_run_step(self, step: "RunStep") -> None:  # type: ignore[func-retur
     async def on_run_step_delta(self, delta: "RunStepDeltaChunk") -> None:  # type: ignore[func-returns-value]
         if self.inner_handler:
             return await self.inner_handler.on_run_step_delta(delta)  # type: ignore
+        return await super().on_run_step_delta(delta)  # type: ignore
 
     async def on_error(self, data: str) -> None:  # type: ignore[func-returns-value]
         if self.inner_handler:
             return await self.inner_handler.on_error(data)  # type: ignore
+        return await super().on_error(data)  # type: ignore
 
     async def on_done(self) -> None:  # type: ignore[func-returns-value]
         if self.inner_handler:
             return await self.inner_handler.on_done()  # type: ignore
+        return await super().on_done()  # type: ignore
         # it could be called multiple tines (for each step) __exit__
 
     async def on_unhandled_event(self, event_type: str, event_data: Any) -> None:  # type: ignore[func-returns-value]
         if self.inner_handler:
             return await self.inner_handler.on_unhandled_event(event_type, event_data)  # type: ignore
+        return await super().on_unhandled_event(event_type, event_data)  # type: ignore
 
     # pylint: enable=R1710