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,23 +504,25 @@ 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"]
-
-            result = AgentResult(*event["stop"])
-            callback_handler(result=result)
-            yield {"result": result}
-
-            self._end_agent_trace_span(response=result)
-
-        except Exception as e:
-            self._end_agent_trace_span(error=e)
-            raise
+        if self.trace_span is not None:
+            span = self.trace_span
+            with trace_api.use_span(span):
+                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"]
+
+                    result = AgentResult(*event["stop"])
+                    callback_handler(result=result)
+                    yield {"result": result}
+
+                    self._end_agent_trace_span(response=result)
+
+                except Exception as e:
+                    self._end_agent_trace_span(error=e)
+                    raise
 
     async def _run_loop(
         self, message: Message, invocation_state: dict[str, Any]
@@ -659,7 +661,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,

src/strands/telemetry/tracer.py

@@ -414,7 +414,7 @@ def start_agent_span(
         tools: Optional[list] = None,
         custom_trace_attributes: Optional[Mapping[str, AttributeValue]] = None,
         **kwargs: Any,
-    ) -> Optional[Span]:
+    ) -> Span:
         """Start a new span for an agent invocation.
 
         Args: