fix(telemetry): group traces when using agent as tool in an agent

Author: poshinchen

src/strands/agent/agent.py

@@ -16,7 +16,7 @@
 from concurrent.futures import ThreadPoolExecutor
 from typing import Any, AsyncGenerator, AsyncIterator, Callable, Mapping, Optional, Type, TypeVar, Union, cast
 
-from opentelemetry import trace
+from opentelemetry import trace as trace_api
 from pydantic import BaseModel
 
 from ..event_loop.event_loop import event_loop_cycle, run_tool
@@ -300,7 +300,7 @@ def __init__(
 
         # Initialize tracer instance (no-op if not configured)
         self.tracer = get_tracer()
-        self.trace_span: Optional[trace.Span] = None
+        self.trace_span: Optional[trace_api.Span] = None
 
         # Initialize agent state management
         if state is not None:
@@ -504,17 +504,17 @@ async def stream_async(self, prompt: Union[str, list[ContentBlock]], **kwargs: A
         message: Message = {"role": "user", "content": content}
 
         self._start_agent_trace_span(message)
-
         try:
-            events = self._run_loop(message, invocation_state=kwargs)
-            async for event in events:
-                if "callback" in event:
-                    callback_handler(**event["callback"])
-                    yield event["callback"]
+            with trace_api.use_span(self.trace_span):
+                events = self._run_loop(message, invocation_state=kwargs)
+                async for event in events:
+                    if "callback" in event:
+                        callback_handler(**event["callback"])
+                        yield event["callback"]
 
-            result = AgentResult(*event["stop"])
-            callback_handler(result=result)
-            yield {"result": result}
+                result = AgentResult(*event["stop"])
+                callback_handler(result=result)
+                yield {"result": result}
 
             self._end_agent_trace_span(response=result)
 
@@ -659,7 +659,6 @@ def _start_agent_trace_span(self, message: Message) -> None:
             message: The user message.
         """
         model_id = self.model.config.get("model_id") if hasattr(self.model, "config") else None
-
         self.trace_span = self.tracer.start_agent_span(
             message=message,
             agent_name=self.name,