chore: allow custom tracer provider to Agent

Author: poshinchen

src/strands/agent/agent.py

@@ -20,6 +20,7 @@
 from uuid import uuid4
 
 from opentelemetry import trace
+from opentelemetry.sdk.trace import TracerProvider
 
 from ..event_loop.event_loop import event_loop_cycle
 from ..handlers.callback_handler import CompositeCallbackHandler, PrintingCallbackHandler, null_callback_handler
@@ -216,6 +217,7 @@ def __init__(
         record_direct_tool_call: bool = True,
         load_tools_from_directory: bool = True,
         trace_attributes: Optional[Mapping[str, AttributeValue]] = None,
+        tracer_provider: Optional[TracerProvider] = None,
     ):
         """Initialize the Agent with the specified configuration.
 
@@ -248,6 +250,7 @@ def __init__(
             load_tools_from_directory: Whether to load and automatically reload tools in the `./tools/` directory.
                 Defaults to True.
             trace_attributes: Custom trace attributes to apply to the agent's trace span.
+            tracer_provider: Custom trace provider to apply to the agents' tracer.
 
         Raises:
             ValueError: If max_parallel_tools is less than 1.
@@ -306,7 +309,7 @@ def __init__(
         self.event_loop_metrics = EventLoopMetrics()
 
         # Initialize tracer instance (no-op if not configured)
-        self.tracer = get_tracer()
+        self.tracer = get_tracer(tracer_provider=tracer_provider)
         self.trace_span: Optional[trace.Span] = None
 
         self.tool_caller = Agent.ToolCaller(self)

src/strands/agent/conversation_manager/sliding_window_conversation_manager.py

@@ -140,7 +140,7 @@ def reduce_context(self, agent: "Agent", e: Optional[Exception] = None) -> None:
             if results_truncated:
                 logger.debug("message_index=<%s> | tool results truncated", last_message_idx_with_tool_results)
                 return
-        
+
         # Try to trim index id when tool result cannot be truncated anymore
         # If the number of messages is less than the window_size, then we default to 2, otherwise, trim to window size
         trim_index = 2 if len(messages) <= self.window_size else len(messages) - self.window_size

src/strands/telemetry/tracer.py

@@ -11,12 +11,16 @@
 from importlib.metadata import version
 from typing import Any, Dict, Mapping, Optional
 
-from opentelemetry import trace
+from opentelemetry import propagate
+from opentelemetry import trace as trace_api
+from opentelemetry.baggage.propagation import W3CBaggagePropagator
 from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
+from opentelemetry.propagators.composite import CompositePropagator
 from opentelemetry.sdk.resources import Resource
-from opentelemetry.sdk.trace import TracerProvider
+from opentelemetry.sdk.trace import TracerProvider as SDKTracerProvider
 from opentelemetry.sdk.trace.export import BatchSpanProcessor, ConsoleSpanExporter, SimpleSpanProcessor
-from opentelemetry.trace import StatusCode
+from opentelemetry.trace import Span, StatusCode
+from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
 
 from ..agent.agent_result import AgentResult
 from ..types.content import Message, Messages
@@ -94,6 +98,7 @@ def __init__(
         otlp_endpoint: Optional[str] = None,
         otlp_headers: Optional[Dict[str, str]] = None,
         enable_console_export: Optional[bool] = None,
+        tracer_provider: Optional[trace_api.TracerProvider] = None,
     ):
         """Initialize the tracer.
 
@@ -102,6 +107,7 @@ def __init__(
             otlp_endpoint: OTLP endpoint URL for sending traces.
             otlp_headers: Headers to include with OTLP requests.
             enable_console_export: Whether to also export traces to console.
+            tracer_provider: Optional existing TracerProvider to use instead of creating a new one.
         """
         # Check environment variables first
         env_endpoint = os.environ.get("OTEL_EXPORTER_OTLP_ENDPOINT")
@@ -133,16 +139,32 @@ def __init__(
 
         self.service_name = service_name
         self.otlp_headers = otlp_headers or {}
-        self.tracer_provider: Optional[TracerProvider] = None
-        self.tracer: Optional[trace.Tracer] = None
-
-        if self.otlp_endpoint or self.enable_console_export:
+        self.tracer_provider = tracer_provider
+        self.tracer: Optional[trace_api.Tracer] = None
+
+        propagate.set_global_textmap(
+            CompositePropagator(
+                [
+                    W3CBaggagePropagator(),
+                    TraceContextTextMapPropagator(),
+                ]
+            )
+        )
+        if self.tracer_provider:
+            # Use the provided tracer provider directly
+            self.tracer = self.tracer_provider.get_tracer(self.service_name)
+        elif self.otlp_endpoint or self.enable_console_export:
+            # Create our own tracer provider
             self._initialize_tracer()
 
     def _initialize_tracer(self) -> None:
         """Initialize the OpenTelemetry tracer."""
         logger.info("initializing tracer")
 
+        # Use global tracer provider if already eixsts
+        if self._use_global_tracer_provider():
+            return
+
         # Create resource with service information
         resource = Resource.create(
             {
@@ -154,7 +176,7 @@ def _initialize_tracer(self) -> None:
         )
 
         # Create tracer provider
-        self.tracer_provider = TracerProvider(resource=resource)
+        self.tracer_provider = SDKTracerProvider(resource=resource)
 
         # Add console exporter if enabled
         if self.enable_console_export and self.tracer_provider:
@@ -190,15 +212,24 @@ def _initialize_tracer(self) -> None:
                 logger.exception("error=<%s> | Failed to configure OTLP exporter", e)
 
         # Set as global tracer provider
-        trace.set_tracer_provider(self.tracer_provider)
-        self.tracer = trace.get_tracer(self.service_name)
+        trace_api.set_tracer_provider(self.tracer_provider)
+        self.tracer = trace_api.get_tracer(self.service_name)
+
+    def _use_global_tracer_provider(self) -> bool:
+        current_provider = trace_api.get_tracer_provider()
+        if isinstance(current_provider, SDKTracerProvider):
+            logger.info("using existing global tracer provider")
+            self.tracer_provider = current_provider
+            self.tracer = current_provider.get_tracer(self.service_name)
+            return True
+        return False
 
     def _start_span(
         self,
         span_name: str,
-        parent_span: Optional[trace.Span] = None,
+        parent_span: Optional[Span] = None,
         attributes: Optional[Dict[str, AttributeValue]] = None,
-    ) -> Optional[trace.Span]:
+    ) -> Optional[Span]:
         """Generic helper method to start a span with common attributes.
 
         Args:
@@ -212,7 +243,7 @@ def _start_span(
         if self.tracer is None:
             return None
 
-        context = trace.set_span_in_context(parent_span) if parent_span else None
+        context = trace_api.set_span_in_context(parent_span) if parent_span else None
         span = self.tracer.start_span(name=span_name, context=context)
 
         # Set start time as a common attribute
@@ -224,7 +255,7 @@ def _start_span(
 
         return span
 
-    def _set_attributes(self, span: trace.Span, attributes: Dict[str, AttributeValue]) -> None:
+    def _set_attributes(self, span: Span, attributes: Dict[str, AttributeValue]) -> None:
         """Set attributes on a span, handling different value types appropriately.
 
         Args:
@@ -239,7 +270,7 @@ def _set_attributes(self, span: trace.Span, attributes: Dict[str, AttributeValue
 
     def _end_span(
         self,
-        span: trace.Span,
+        span: Span,
         attributes: Optional[Dict[str, AttributeValue]] = None,
         error: Optional[Exception] = None,
     ) -> None:
@@ -278,7 +309,7 @@ def _end_span(
                 except Exception as e:
                     logger.warning("error=<%s> | failed to force flush tracer provider", e)
 
-    def end_span_with_error(self, span: trace.Span, error_message: str, exception: Optional[Exception] = None) -> None:
+    def end_span_with_error(self, span: Span, error_message: str, exception: Optional[Exception] = None) -> None:
         """End a span with error status.
 
         Args:
@@ -294,12 +325,12 @@ def end_span_with_error(self, span: trace.Span, error_message: str, exception: O
 
     def start_model_invoke_span(
         self,
-        parent_span: Optional[trace.Span] = None,
+        parent_span: Optional[Span] = None,
         agent_name: str = "Strands Agent",
         messages: Optional[Messages] = None,
         model_id: Optional[str] = None,
         **kwargs: Any,
-    ) -> Optional[trace.Span]:
+    ) -> Optional[Span]:
         """Start a new span for a model invocation.
 
         Args:
@@ -328,7 +359,7 @@ def start_model_invoke_span(
         return self._start_span("Model invoke", parent_span, attributes)
 
     def end_model_invoke_span(
-        self, span: trace.Span, message: Message, usage: Usage, error: Optional[Exception] = None
+        self, span: Span, message: Message, usage: Usage, error: Optional[Exception] = None
     ) -> None:
         """End a model invocation span with results and metrics.
 
@@ -347,9 +378,7 @@ def end_model_invoke_span(
 
         self._end_span(span, attributes, error)
 
-    def start_tool_call_span(
-        self, tool: ToolUse, parent_span: Optional[trace.Span] = None, **kwargs: Any
-    ) -> Optional[trace.Span]:
+    def start_tool_call_span(self, tool: ToolUse, parent_span: Optional[Span] = None, **kwargs: Any) -> Optional[Span]:
         """Start a new span for a tool call.
 
         Args:
@@ -374,7 +403,7 @@ def start_tool_call_span(
         return self._start_span(span_name, parent_span, attributes)
 
     def end_tool_call_span(
-        self, span: trace.Span, tool_result: Optional[ToolResult], error: Optional[Exception] = None
+        self, span: Span, tool_result: Optional[ToolResult], error: Optional[Exception] = None
     ) -> None:
         """End a tool call span with results.
 
@@ -402,10 +431,10 @@ def end_tool_call_span(
     def start_event_loop_cycle_span(
         self,
         event_loop_kwargs: Any,
-        parent_span: Optional[trace.Span] = None,
+        parent_span: Optional[Span] = None,
         messages: Optional[Messages] = None,
         **kwargs: Any,
-    ) -> Optional[trace.Span]:
+    ) -> Optional[Span]:
         """Start a new span for an event loop cycle.
 
         Args:
@@ -436,7 +465,7 @@ def start_event_loop_cycle_span(
 
     def end_event_loop_cycle_span(
         self,
-        span: trace.Span,
+        span: Span,
         message: Message,
         tool_result_message: Optional[Message] = None,
         error: Optional[Exception] = None,
@@ -466,7 +495,7 @@ def start_agent_span(
         tools: Optional[list] = None,
         custom_trace_attributes: Optional[Mapping[str, AttributeValue]] = None,
         **kwargs: Any,
-    ) -> Optional[trace.Span]:
+    ) -> Optional[Span]:
         """Start a new span for an agent invocation.
 
         Args:
@@ -506,7 +535,7 @@ def start_agent_span(
 
     def end_agent_span(
         self,
-        span: trace.Span,
+        span: Span,
         response: Optional[AgentResult] = None,
         error: Optional[Exception] = None,
     ) -> None:
@@ -549,6 +578,7 @@ def get_tracer(
     otlp_endpoint: Optional[str] = None,
     otlp_headers: Optional[Dict[str, str]] = None,
     enable_console_export: Optional[bool] = None,
+    tracer_provider: Optional[trace_api.TracerProvider] = None,
 ) -> Tracer:
     """Get or create the global tracer.
 
@@ -557,18 +587,22 @@ def get_tracer(
         otlp_endpoint: OTLP endpoint URL for sending traces.
         otlp_headers: Headers to include with OTLP requests.
         enable_console_export: Whether to also export traces to console.
+        tracer_provider: Optional existing TracerProvider to use instead of creating a new one.
 
     Returns:
         The global tracer instance.
     """
     global _tracer_instance
 
-    if _tracer_instance is None or (otlp_endpoint and _tracer_instance.otlp_endpoint != otlp_endpoint):  # type: ignore[unreachable]
+    if (
+        _tracer_instance is None or (otlp_endpoint and _tracer_instance.otlp_endpoint != otlp_endpoint)  # type: ignore[unreachable]
+    ):
         _tracer_instance = Tracer(
             service_name=service_name,
             otlp_endpoint=otlp_endpoint,
             otlp_headers=otlp_headers,
             enable_console_export=enable_console_export,
+            tracer_provider=tracer_provider,
         )
 
     return _tracer_instance