fix: include agent trace into tool for agent as tools (#526)

Author: poshinchen

src/strands/telemetry/tracer.py

@@ -273,7 +273,7 @@ def end_model_invoke_span(
 
         self._end_span(span, attributes, error)
 
-    def start_tool_call_span(self, tool: ToolUse, parent_span: Optional[Span] = None, **kwargs: Any) -> Optional[Span]:
+    def start_tool_call_span(self, tool: ToolUse, parent_span: Optional[Span] = None, **kwargs: Any) -> Span:
         """Start a new span for a tool call.
 
         Args:

src/strands/tools/executor.py

@@ -5,7 +5,7 @@
 import time
 from typing import Any, Optional, cast
 
-from opentelemetry import trace
+from opentelemetry import trace as trace_api
 
 from ..telemetry.metrics import EventLoopMetrics, Trace
 from ..telemetry.tracer import get_tracer
@@ -23,7 +23,7 @@ async def run_tools(
     invalid_tool_use_ids: list[str],
     tool_results: list[ToolResult],
     cycle_trace: Trace,
-    parent_span: Optional[trace.Span] = None,
+    parent_span: Optional[trace_api.Span] = None,
 ) -> ToolGenerator:
     """Execute tools concurrently.
 
@@ -53,24 +53,23 @@ async def work(
         tool_name = tool_use["name"]
         tool_trace = Trace(f"Tool: {tool_name}", parent_id=cycle_trace.id, raw_name=tool_name)
         tool_start_time = time.time()
+        with trace_api.use_span(tool_call_span):
+            try:
+                async for event in handler(tool_use):
+                    worker_queue.put_nowait((worker_id, event))
+                    await worker_event.wait()
+                    worker_event.clear()
+
+                result = cast(ToolResult, event)
+            finally:
+                worker_queue.put_nowait((worker_id, stop_event))
+
+            tool_success = result.get("status") == "success"
+            tool_duration = time.time() - tool_start_time
+            message = Message(role="user", content=[{"toolResult": result}])
+            event_loop_metrics.add_tool_usage(tool_use, tool_duration, tool_trace, tool_success, message)
+            cycle_trace.add_child(tool_trace)
 
-        try:
-            async for event in handler(tool_use):
-                worker_queue.put_nowait((worker_id, event))
-                await worker_event.wait()
-                worker_event.clear()
-
-            result = cast(ToolResult, event)
-        finally:
-            worker_queue.put_nowait((worker_id, stop_event))
-
-        tool_success = result.get("status") == "success"
-        tool_duration = time.time() - tool_start_time
-        message = Message(role="user", content=[{"toolResult": result}])
-        event_loop_metrics.add_tool_usage(tool_use, tool_duration, tool_trace, tool_success, message)
-        cycle_trace.add_child(tool_trace)
-
-        if tool_call_span:
             tracer.end_tool_call_span(tool_call_span, result)
 
         return result