fix(tracing): use a span processor to add custom span attributes (#1104)

- Fix custom tracing span attributes for the kagent-adk,
kagent-langgraph, kagent-crewai packages.
- Add a custom span processor which can be used to include custom kagent
span metrics for all spans created within a particular OTel trace
context.
- Use the custom span processor, and update the execute method of each
agent framework package to set the custom span attributes in context.

Signed-off-by: David Haifley <[email protected]>

Author: dhaifley

python/packages/kagent-adk/src/kagent/adk/_agent_executor.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import asyncio
 import inspect
 import logging
 import uuid
@@ -29,6 +28,10 @@
 from typing_extensions import override
 
 from kagent.core.a2a import TaskResultAggregator, get_kagent_metadata_key
+from kagent.core.tracing._span_processor import (
+    clear_kagent_span_attributes,
+    set_kagent_span_attributes,
+)
 
 from .converters.event_converter import convert_event_to_a2a_events
 from .converters.request_converter import convert_a2a_request_to_adk_run_args
@@ -106,57 +109,72 @@ async def execute(
         if not context.message:
             raise ValueError("A2A request must have a message")
 
-        # for new task, create a task submitted event
-        if not context.current_task:
-            await event_queue.enqueue_event(
-                TaskStatusUpdateEvent(
-                    task_id=context.task_id,
-                    status=TaskStatus(
-                        state=TaskState.submitted,
-                        message=context.message,
-                        timestamp=datetime.now(timezone.utc).isoformat(),
-                    ),
-                    context_id=context.context_id,
-                    final=False,
-                )
-            )
+        # Convert the a2a request to ADK run args
+        run_args = convert_a2a_request_to_adk_run_args(context)
 
-        # Handle the request and publish updates to the event queue
-        runner = await self._resolve_runner()
+        # Prepare span attributes.
+        span_attributes = {}
+        if run_args.get("user_id"):
+            span_attributes["kagent.user_id"] = run_args["user_id"]
+        if context.task_id:
+            span_attributes["gen_ai.task.id"] = context.task_id
+        if run_args.get("session_id"):
+            span_attributes["gen_ai.conversation.id"] = run_args["session_id"]
+
+        # Set kagent span attributes for all spans in context.
+        context_token = set_kagent_span_attributes(span_attributes)
         try:
-            await self._handle_request(context, event_queue, runner)
-        except Exception as e:
-            logger.error("Error handling A2A request: %s", e, exc_info=True)
-            # Publish failure event
-            try:
+            # for new task, create a task submitted event
+            if not context.current_task:
                 await event_queue.enqueue_event(
                     TaskStatusUpdateEvent(
                         task_id=context.task_id,
                         status=TaskStatus(
-                            state=TaskState.failed,
+                            state=TaskState.submitted,
+                            message=context.message,
                             timestamp=datetime.now(timezone.utc).isoformat(),
-                            message=Message(
-                                message_id=str(uuid.uuid4()),
-                                role=Role.agent,
-                                parts=[Part(TextPart(text=str(e)))],
-                            ),
                         ),
                         context_id=context.context_id,
-                        final=True,
+                        final=False,
                     )
                 )
-            except Exception as enqueue_error:
-                logger.error("Failed to publish failure event: %s", enqueue_error, exc_info=True)
+
+            # Handle the request and publish updates to the event queue
+            runner = await self._resolve_runner()
+            try:
+                await self._handle_request(context, event_queue, runner, run_args)
+            except Exception as e:
+                logger.error("Error handling A2A request: %s", e, exc_info=True)
+                # Publish failure event
+                try:
+                    await event_queue.enqueue_event(
+                        TaskStatusUpdateEvent(
+                            task_id=context.task_id,
+                            status=TaskStatus(
+                                state=TaskState.failed,
+                                timestamp=datetime.now(timezone.utc).isoformat(),
+                                message=Message(
+                                    message_id=str(uuid.uuid4()),
+                                    role=Role.agent,
+                                    parts=[Part(TextPart(text=str(e)))],
+                                ),
+                            ),
+                            context_id=context.context_id,
+                            final=True,
+                        )
+                    )
+                except Exception as enqueue_error:
+                    logger.error("Failed to publish failure event: %s", enqueue_error, exc_info=True)
+        finally:
+            clear_kagent_span_attributes(context_token)
 
     async def _handle_request(
         self,
         context: RequestContext,
         event_queue: EventQueue,
         runner: Runner,
+        run_args: dict[str, Any],
     ):
-        # Convert the a2a request to ADK run args
-        run_args = convert_a2a_request_to_adk_run_args(context)
-
         # ensure the session exists
         session = await self._prepare_session(context, run_args, runner)
 
@@ -175,14 +193,6 @@ async def _handle_request(
 
         await runner.session_service.append_event(session, system_event)
 
-        current_span = trace.get_current_span()
-        if run_args["user_id"]:
-            current_span.set_attribute("kagent.user_id", run_args["user_id"])
-        if context.task_id:
-            current_span.set_attribute("gen_ai.task.id", context.task_id)
-        if run_args["session_id"]:
-            current_span.set_attribute("gen_ai.converstation.id", run_args["session_id"])
-
         # create invocation context
         invocation_context = runner._new_invocation_context(
             session=session,
@@ -236,7 +246,7 @@ async def _handle_request(
                     ),
                 )
             )
-            # public the final status update event
+            # publish the final status update event
             await event_queue.enqueue_event(
                 TaskStatusUpdateEvent(
                     task_id=context.task_id,

python/packages/kagent-core/src/kagent/core/tracing/_span_processor.py

@@ -0,0 +1,59 @@
+"""Custom span processor to add kagent attributes to all spans in a request context."""
+
+import logging
+from typing import Optional
+from contextvars import Token
+
+from opentelemetry import context as otel_context
+from opentelemetry.sdk.trace import ReadableSpan, Span, SpanProcessor
+
+logger = logging.getLogger(__name__)
+
+KAGENT_ATTRIBUTES_KEY = "kagent_trace_span_attributes"
+
+
+class KagentAttributesSpanProcessor(SpanProcessor):
+    """A SpanProcessor that adds kagent-specific attributes to all spans."""
+
+    def on_start(self, span: Span, parent_context: Optional[otel_context.Context] = None) -> None:
+        """Called when a span is started. Adds kagent attributes if present in context."""
+        try:
+            ctx = parent_context if parent_context is not None else otel_context.get_current()
+            attributes = ctx.get(KAGENT_ATTRIBUTES_KEY)
+
+            if attributes and isinstance(attributes, dict):
+                for key, value in attributes.items():
+                    if value is not None:
+                        span.set_attribute(key, value)
+        except Exception as e:
+            logger.warning(f"Failed to add kagent attributes to span: {e}")
+
+    def on_end(self, span: ReadableSpan) -> None:
+        """Called when a span is ended. No action needed."""
+        pass
+
+    def shutdown(self) -> None:
+        """Called when the tracer provider is shutdown."""
+        pass
+
+    def force_flush(self, timeout_millis: int = 30000) -> bool:
+        """Force flush any buffered spans. No buffering in this processor."""
+        return True
+
+
+def set_kagent_span_attributes(attributes: dict) -> Token[otel_context.Context]:
+    """Set kagent span attributes in the context.
+    Args:
+        attributes: Dictionary of kagent span attributes to store in context
+    Returns:
+        A context token that can be used to detach the context
+    """
+    return otel_context.attach(otel_context.set_value(KAGENT_ATTRIBUTES_KEY, attributes))
+
+
+def clear_kagent_span_attributes(token: Token[otel_context.Context]) -> None:
+    """Clear kagent span attributes from the OpenTelemetry context.
+    Args:
+        token: The context token returned by set_kagent_span_attributes
+    """
+    otel_context.detach(token)

python/packages/kagent-core/src/kagent/core/tracing/_utils.py

@@ -17,6 +17,8 @@
 from opentelemetry.sdk.trace import TracerProvider
 from opentelemetry.sdk.trace.export import BatchSpanProcessor
 
+from ._span_processor import KagentAttributesSpanProcessor
+
 
 def configure(fastapi_app: FastAPI | None = None):
     tracing_enabled = os.getenv("OTEL_TRACING_ENABLED", "false").lower() == "true"
@@ -38,13 +40,15 @@ def configure(fastapi_app: FastAPI | None = None):
         # Check if a TracerProvider already exists (e.g., set by CrewAI)
         current_provider = trace.get_tracer_provider()
         if isinstance(current_provider, TracerProvider):
-            # TracerProvider already exists, just add our processor to it
+            # TracerProvider already exists, just add our processors to it
             current_provider.add_span_processor(processor)
-            logging.info("Added OTLP processor to existing TracerProvider")
+            current_provider.add_span_processor(KagentAttributesSpanProcessor())
+            logging.info("Added OTLP processors to existing TracerProvider")
         else:
             # No provider set, create new one
             tracer_provider = TracerProvider(resource=resource)
             tracer_provider.add_span_processor(processor)
+            tracer_provider.add_span_processor(KagentAttributesSpanProcessor())
             trace.set_tracer_provider(tracer_provider)
             logging.info("Created new TracerProvider")