fix: attached custom attributes to all spans (#1235)

Author: poshinchen

src/strands/event_loop/event_loop.py

@@ -133,7 +133,10 @@ async def event_loop_cycle(
     # Create tracer span for this event loop cycle
     tracer = get_tracer()
     cycle_span = tracer.start_event_loop_cycle_span(
-        invocation_state=invocation_state, messages=agent.messages, parent_span=agent.trace_span
+        invocation_state=invocation_state,
+        messages=agent.messages,
+        parent_span=agent.trace_span,
+        custom_trace_attributes=agent.trace_attributes,
     )
     invocation_state["event_loop_cycle_span"] = cycle_span
 
@@ -320,6 +323,7 @@ async def _handle_model_execution(
             messages=agent.messages,
             parent_span=cycle_span,
             model_id=model_id,
+            custom_trace_attributes=agent.trace_attributes,
         )
         with trace_api.use_span(model_invoke_span):
             await agent.hooks.invoke_callbacks_async(

src/strands/multiagent/base.py

@@ -8,12 +8,13 @@
 from abc import ABC, abstractmethod
 from dataclasses import dataclass, field
 from enum import Enum
-from typing import Any, AsyncIterator, Union
+from typing import Any, AsyncIterator, Mapping, Union
 
 from .._async import run_async
 from ..agent import AgentResult
 from ..types.event_loop import Metrics, Usage
 from ..types.multiagent import MultiAgentInput
+from ..types.traces import AttributeValue
 
 logger = logging.getLogger(__name__)
 
@@ -238,6 +239,18 @@ def deserialize_state(self, payload: dict[str, Any]) -> None:
         """Restore orchestrator state from a session dict."""
         raise NotImplementedError
 
+    def _parse_trace_attributes(
+        self, attributes: Mapping[str, AttributeValue] | None = None
+    ) -> dict[str, AttributeValue]:
+        trace_attributes: dict[str, AttributeValue] = {}
+        if attributes:
+            for k, v in attributes.items():
+                if isinstance(v, (str, int, float, bool)) or (
+                    isinstance(v, list) and all(isinstance(x, (str, int, float, bool)) for x in v)
+                ):
+                    trace_attributes[k] = v
+        return trace_attributes
+
 
 # Private helper function to avoid duplicate code
 

src/strands/multiagent/graph.py

@@ -19,7 +19,7 @@
 import logging
 import time
 from dataclasses import dataclass, field
-from typing import Any, AsyncIterator, Callable, Optional, Tuple, cast
+from typing import Any, AsyncIterator, Callable, Mapping, Optional, Tuple, cast
 
 from opentelemetry import trace as trace_api
 
@@ -46,6 +46,7 @@
 from ..types.content import ContentBlock, Messages
 from ..types.event_loop import Metrics, Usage
 from ..types.multiagent import MultiAgentInput
+from ..types.traces import AttributeValue
 from .base import MultiAgentBase, MultiAgentResult, NodeResult, Status
 
 logger = logging.getLogger(__name__)
@@ -413,6 +414,7 @@ def __init__(
         session_manager: Optional[SessionManager] = None,
         hooks: Optional[list[HookProvider]] = None,
         id: str = _DEFAULT_GRAPH_ID,
+        trace_attributes: Optional[Mapping[str, AttributeValue]] = None,
     ) -> None:
         """Initialize Graph with execution limits and reset behavior.
 
@@ -427,6 +429,7 @@ def __init__(
             session_manager: Session manager for persisting graph state and execution history (default: None)
             hooks: List of hook providers for monitoring and extending graph execution behavior (default: None)
             id: Unique graph id (default: None)
+            trace_attributes: Custom trace attributes to apply to the agent's trace span (default: None)
         """
         super().__init__()
 
@@ -442,6 +445,7 @@ def __init__(
         self.reset_on_revisit = reset_on_revisit
         self.state = GraphState()
         self.tracer = get_tracer()
+        self.trace_attributes: dict[str, AttributeValue] = self._parse_trace_attributes(trace_attributes)
         self.session_manager = session_manager
         self.hooks = HookRegistry()
         if self.session_manager:
@@ -537,7 +541,7 @@ async def stream_async(
             self.state.status = Status.EXECUTING
             self.state.start_time = start_time
 
-        span = self.tracer.start_multiagent_span(task, "graph")
+        span = self.tracer.start_multiagent_span(task, "graph", custom_trace_attributes=self.trace_attributes)
         with trace_api.use_span(span, end_on_exit=True):
             try:
                 logger.debug(

src/strands/multiagent/swarm.py

@@ -18,7 +18,7 @@
 import logging
 import time
 from dataclasses import dataclass, field
-from typing import Any, AsyncIterator, Callable, Optional, Tuple, cast
+from typing import Any, AsyncIterator, Callable, Mapping, Optional, Tuple, cast
 
 from opentelemetry import trace as trace_api
 
@@ -46,6 +46,7 @@
 from ..types.content import ContentBlock, Messages
 from ..types.event_loop import Metrics, Usage
 from ..types.multiagent import MultiAgentInput
+from ..types.traces import AttributeValue
 from .base import MultiAgentBase, MultiAgentResult, NodeResult, Status
 
 logger = logging.getLogger(__name__)
@@ -226,6 +227,7 @@ def __init__(
         session_manager: Optional[SessionManager] = None,
         hooks: Optional[list[HookProvider]] = None,
         id: str = _DEFAULT_SWARM_ID,
+        trace_attributes: Optional[Mapping[str, AttributeValue]] = None,
     ) -> None:
         """Initialize Swarm with agents and configuration.
 
@@ -243,6 +245,7 @@ def __init__(
                 Disabled by default (default: 0)
             session_manager: Session manager for persisting graph state and execution history (default: None)
             hooks: List of hook providers for monitoring and extending graph execution behavior (default: None)
+            trace_attributes: Custom trace attributes to apply to the agent's trace span (default: None)
         """
         super().__init__()
         self.id = id
@@ -262,6 +265,7 @@ def __init__(
             completion_status=Status.PENDING,
         )
         self.tracer = get_tracer()
+        self.trace_attributes: dict[str, AttributeValue] = self._parse_trace_attributes(trace_attributes)
 
         self.session_manager = session_manager
         self.hooks = HookRegistry()
@@ -356,7 +360,7 @@ async def stream_async(
             self.state.completion_status = Status.EXECUTING
             self.state.start_time = time.time()
 
-        span = self.tracer.start_multiagent_span(task, "swarm")
+        span = self.tracer.start_multiagent_span(task, "swarm", custom_trace_attributes=self.trace_attributes)
         with trace_api.use_span(span, end_on_exit=True):
             try:
                 current_node = cast(SwarmNode, self.state.current_node)

src/strands/telemetry/tracer.py

@@ -277,6 +277,7 @@ def start_model_invoke_span(
         messages: Messages,
         parent_span: Optional[Span] = None,
         model_id: Optional[str] = None,
+        custom_trace_attributes: Optional[Mapping[str, AttributeValue]] = None,
         **kwargs: Any,
     ) -> Span:
         """Start a new span for a model invocation.
@@ -285,13 +286,17 @@ def start_model_invoke_span(
             messages: Messages being sent to the model.
             parent_span: Optional parent span to link this span to.
             model_id: Optional identifier for the model being invoked.
+            custom_trace_attributes: Optional mapping of custom trace attributes to include in the span.
             **kwargs: Additional attributes to add to the span.
 
         Returns:
             The created span, or None if tracing is not enabled.
         """
         attributes: Dict[str, AttributeValue] = self._get_common_attributes(operation_name="chat")
 
+        if custom_trace_attributes:
+            attributes.update(custom_trace_attributes)
+
         if model_id:
             attributes["gen_ai.request.model"] = model_id
 
@@ -358,12 +363,19 @@ 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) -> Span:
+    def start_tool_call_span(
+        self,
+        tool: ToolUse,
+        parent_span: Optional[Span] = None,
+        custom_trace_attributes: Optional[Mapping[str, AttributeValue]] = None,
+        **kwargs: Any,
+    ) -> Span:
         """Start a new span for a tool call.
 
         Args:
             tool: The tool being used.
             parent_span: Optional parent span to link this span to.
+            custom_trace_attributes: Optional mapping of custom trace attributes to include in the span.
             **kwargs: Additional attributes to add to the span.
 
         Returns:
@@ -377,6 +389,8 @@ def start_tool_call_span(self, tool: ToolUse, parent_span: Optional[Span] = None
             }
         )
 
+        if custom_trace_attributes:
+            attributes.update(custom_trace_attributes)
         # Add additional kwargs as attributes
         attributes.update(kwargs)
 
@@ -477,6 +491,7 @@ def start_event_loop_cycle_span(
         invocation_state: Any,
         messages: Messages,
         parent_span: Optional[Span] = None,
+        custom_trace_attributes: Optional[Mapping[str, AttributeValue]] = None,
         **kwargs: Any,
     ) -> Optional[Span]:
         """Start a new span for an event loop cycle.
@@ -485,6 +500,7 @@ def start_event_loop_cycle_span(
             invocation_state: Arguments for the event loop cycle.
             parent_span: Optional parent span to link this span to.
             messages:  Messages being processed in this cycle.
+            custom_trace_attributes: Optional mapping of custom trace attributes to include in the span.
             **kwargs: Additional attributes to add to the span.
 
         Returns:
@@ -497,6 +513,9 @@ def start_event_loop_cycle_span(
             "event_loop.cycle_id": event_loop_cycle_id,
         }
 
+        if custom_trace_attributes:
+            attributes.update(custom_trace_attributes)
+
         if "event_loop_parent_cycle_id" in invocation_state:
             attributes["event_loop.parent_cycle_id"] = str(invocation_state["event_loop_parent_cycle_id"])
 
@@ -679,6 +698,7 @@ def start_multiagent_span(
         self,
         task: MultiAgentInput,
         instance: str,
+        custom_trace_attributes: Optional[Mapping[str, AttributeValue]] = None,
     ) -> Span:
         """Start a new span for swarm invocation."""
         operation = f"invoke_{instance}"
@@ -689,6 +709,9 @@ def start_multiagent_span(
             }
         )
 
+        if custom_trace_attributes:
+            attributes.update(custom_trace_attributes)
+
         span = self._start_span(operation, attributes=attributes, span_kind=trace_api.SpanKind.CLIENT)
 
         if self.use_latest_genai_conventions:

src/strands/tools/executors/_executor.py

@@ -249,7 +249,9 @@ async def _stream_with_trace(
 
         tracer = get_tracer()
 
-        tool_call_span = tracer.start_tool_call_span(tool_use, cycle_span)
+        tool_call_span = tracer.start_tool_call_span(
+            tool_use, cycle_span, custom_trace_attributes=agent.trace_attributes
+        )
         tool_trace = Trace(f"Tool: {tool_name}", parent_id=cycle_trace.id, raw_name=tool_name)
         tool_start_time = time.time()
 

tests/strands/event_loop/test_event_loop.py

@@ -143,6 +143,7 @@ def agent(model, system_prompt, messages, tool_registry, thread_pool, hook_regis
     mock.hooks = hook_registry
     mock.tool_executor = tool_executor
     mock._interrupt_state = _InterruptState()
+    mock.trace_attributes = {}
 
     return mock
 
@@ -738,7 +739,10 @@ async def test_event_loop_cycle_with_parent_span(
 
     # Verify parent_span was used when creating cycle span
     mock_tracer.start_event_loop_cycle_span.assert_called_once_with(
-        invocation_state=unittest.mock.ANY, parent_span=parent_span, messages=messages
+        invocation_state=unittest.mock.ANY,
+        parent_span=parent_span,
+        messages=messages,
+        custom_trace_attributes=unittest.mock.ANY,
     )
 
 

tests/strands/event_loop/test_event_loop_structured_output.py

@@ -40,6 +40,7 @@ def mock_agent():
     agent.hooks = Mock()
     agent.hooks.invoke_callbacks_async = AsyncMock()
     agent.trace_span = None
+    agent.trace_attributes = {}
     agent.tool_executor = Mock()
     agent._append_message = AsyncMock()
 

tests/strands/telemetry/test_tracer.py

@@ -149,15 +149,20 @@ def test_start_model_invoke_span(mock_tracer):
 
         messages = [{"role": "user", "content": [{"text": "Hello"}]}]
         model_id = "test-model"
+        custom_attrs = {"custom_key": "custom_value", "user_id": "12345"}
 
-        span = tracer.start_model_invoke_span(messages=messages, agent_name="TestAgent", model_id=model_id)
+        span = tracer.start_model_invoke_span(
+            messages=messages, agent_name="TestAgent", model_id=model_id, custom_trace_attributes=custom_attrs
+        )
 
         mock_tracer.start_span.assert_called_once()
         assert mock_tracer.start_span.call_args[1]["name"] == "chat"
         assert mock_tracer.start_span.call_args[1]["kind"] == SpanKind.INTERNAL
         mock_span.set_attribute.assert_any_call("gen_ai.system", "strands-agents")
         mock_span.set_attribute.assert_any_call("gen_ai.operation.name", "chat")
         mock_span.set_attribute.assert_any_call("gen_ai.request.model", model_id)
+        mock_span.set_attribute.assert_any_call("custom_key", "custom_value")
+        mock_span.set_attribute.assert_any_call("user_id", "12345")
         mock_span.add_event.assert_called_with(
             "gen_ai.user.message", attributes={"content": json.dumps(messages[0]["content"])}
         )
@@ -293,15 +298,18 @@ def test_start_tool_call_span(mock_tracer):
         mock_tracer.start_span.return_value = mock_span
 
         tool = {"name": "test-tool", "toolUseId": "123", "input": {"param": "value"}}
+        custom_attrs = {"session_id": "abc123", "environment": "production"}
 
-        span = tracer.start_tool_call_span(tool)
+        span = tracer.start_tool_call_span(tool, custom_trace_attributes=custom_attrs)
 
         mock_tracer.start_span.assert_called_once()
         assert mock_tracer.start_span.call_args[1]["name"] == "execute_tool test-tool"
         mock_span.set_attribute.assert_any_call("gen_ai.tool.name", "test-tool")
         mock_span.set_attribute.assert_any_call("gen_ai.system", "strands-agents")
         mock_span.set_attribute.assert_any_call("gen_ai.operation.name", "execute_tool")
         mock_span.set_attribute.assert_any_call("gen_ai.tool.call.id", "123")
+        mock_span.set_attribute.assert_any_call("session_id", "abc123")
+        mock_span.set_attribute.assert_any_call("environment", "production")
         mock_span.add_event.assert_any_call(
             "gen_ai.tool.message", attributes={"role": "tool", "content": json.dumps({"param": "value"}), "id": "123"}
         )
@@ -361,14 +369,17 @@ def test_start_swarm_call_span_with_string_task(mock_tracer):
         mock_tracer.start_span.return_value = mock_span
 
         task = "Design foo bar"
+        custom_attrs = {"workflow_id": "wf-789", "priority": "high"}
 
-        span = tracer.start_multiagent_span(task, "swarm")
+        span = tracer.start_multiagent_span(task, "swarm", custom_trace_attributes=custom_attrs)
 
         mock_tracer.start_span.assert_called_once()
         assert mock_tracer.start_span.call_args[1]["name"] == "invoke_swarm"
         mock_span.set_attribute.assert_any_call("gen_ai.system", "strands-agents")
         mock_span.set_attribute.assert_any_call("gen_ai.agent.name", "swarm")
         mock_span.set_attribute.assert_any_call("gen_ai.operation.name", "invoke_swarm")
+        mock_span.set_attribute.assert_any_call("workflow_id", "wf-789")
+        mock_span.set_attribute.assert_any_call("priority", "high")
         mock_span.add_event.assert_any_call("gen_ai.user.message", attributes={"content": "Design foo bar"})
         assert span is not None
 
@@ -575,12 +586,17 @@ def test_start_event_loop_cycle_span(mock_tracer):
 
         event_loop_kwargs = {"event_loop_cycle_id": "cycle-123"}
         messages = [{"role": "user", "content": [{"text": "Hello"}]}]
+        custom_attrs = {"request_id": "req-456", "trace_level": "debug"}
 
-        span = tracer.start_event_loop_cycle_span(event_loop_kwargs, messages=messages)
+        span = tracer.start_event_loop_cycle_span(
+            event_loop_kwargs, messages=messages, custom_trace_attributes=custom_attrs
+        )
 
         mock_tracer.start_span.assert_called_once()
         assert mock_tracer.start_span.call_args[1]["name"] == "execute_event_loop_cycle"
         mock_span.set_attribute.assert_any_call("event_loop.cycle_id", "cycle-123")
+        mock_span.set_attribute.assert_any_call("request_id", "req-456")
+        mock_span.set_attribute.assert_any_call("trace_level", "debug")
         mock_span.add_event.assert_any_call(
             "gen_ai.user.message", attributes={"content": json.dumps([{"text": "Hello"}])}
         )

tests/strands/tools/executors/conftest.py

@@ -105,6 +105,7 @@ def agent(tool_registry, hook_registry):
     mock_agent.tool_registry = tool_registry
     mock_agent.hooks = hook_registry
     mock_agent._interrupt_state = _InterruptState()
+    mock_agent.trace_attributes = {}
     return mock_agent