Telemetry tracing - Enable traces

pyproject.toml

@@ -7,6 +7,7 @@ requires-python = ">=3.13"
 dependencies = [
     "httpx>=0.28.1",
     "httpx-sse>=0.4.0",
+    "mypy>=1.15.0",
     "opentelemetry-api>=1.33.0",
     "opentelemetry-sdk>=1.33.0",
     "pydantic>=2.11.3",

src/a2a/client/client.py

@@ -6,9 +6,9 @@
 
 import httpx
 
 from httpx_sse import SSEError, aconnect_sse
 
 from a2a.client.errors import A2AClientHTTPError, A2AClientJSONError
 from a2a.types import (
     AgentCard,
     CancelTaskRequest,
@@ -24,9 +24,10 @@
     SetTaskPushNotificationConfigRequest,
     SetTaskPushNotificationConfigResponse,
 )
+from a2a.utils.telemetry import SpanKind, trace_class
 
 
 class A2ACardResolver:
     """Agent Card resolver."""
 
     def __init__(
@@ -50,15 +51,16 @@
             response.raise_for_status()
             return AgentCard.model_validate(response.json())
         except httpx.HTTPStatusError as e:
             raise A2AClientHTTPError(e.response.status_code, str(e)) from e
         except json.JSONDecodeError as e:
             raise A2AClientJSONError(str(e)) from e
         except httpx.RequestError as e:
             raise A2AClientHTTPError(
                 503, f'Network communication error: {e}'
             ) from e
 
 
+@trace_class(kind=SpanKind.CLIENT)
 class A2AClient:
     """A2A Client."""
 
@@ -85,7 +87,7 @@
         http_kwargs: dict[str, Any] | None = None,
     ) -> 'A2AClient':
         """Get a A2A client for provided agent card URL."""
         agent_card: AgentCard = await A2ACardResolver(
             httpx_client, base_url=base_url, agent_card_path=agent_card_path
         ).get_agent_card(http_kwargs=http_kwargs)
         return A2AClient(httpx_client=httpx_client, agent_card=agent_card)
@@ -121,7 +123,7 @@
             **(http_kwargs or {}),
         }
 
         async with aconnect_sse(
             self.httpx_client,
             'POST',
             self.url,

src/a2a/server/events/event_consumer.py

@@ -12,11 +12,13 @@
     TaskStatusUpdateEvent,
 )
 from a2a.utils.errors import ServerError
+from a2a.utils.telemetry import SpanKind, trace_class
 
 
 logger = logging.getLogger(__name__)
 
 
+@trace_class(kind=SpanKind.SERVER)
 class EventConsumer:
     """Consumer to read events from the agent event queue."""
 

src/a2a/server/events/event_queue.py

@@ -11,6 +11,7 @@
     TaskArtifactUpdateEvent,
     TaskStatusUpdateEvent,
 )
+from a2a.utils.telemetry import SpanKind, trace_class
 
 
 logger = logging.getLogger(__name__)
@@ -26,6 +27,7 @@
 )
 
 
+@trace_class(kind=SpanKind.SERVER)
 class EventQueue:
     """Event queue for A2A responses from agent."""
 

src/a2a/server/events/in_memory_queue_manager.py

@@ -6,8 +6,10 @@
     QueueManager,
     TaskQueueExists,
 )
+from a2a.utils.telemetry import SpanKind, trace_class
 
 
+@trace_class(kind=SpanKind.SERVER)
 class InMemoryQueueManager(QueueManager):
     """InMemoryQueueManager is used for a single binary management.
 

src/a2a/server/request_handlers/default_request_handler.py

@@ -28,10 +28,12 @@
     UnsupportedOperationError,
 )
 from a2a.utils.errors import ServerError
+from a2a.utils.telemetry import SpanKind, trace_class
 
 logger = logging.getLogger(__name__)
 
 
+@trace_class(kind=SpanKind.SERVER)
 class DefaultRequestHandler(RequestHandler):
     """Default request handler for all incoming requests."""
 
@@ -128,7 +130,7 @@ async def on_message_send(
         # agents.
         queue = await self._queue_manager.create_or_tap(task_id)
         result_aggregator = ResultAggregator(task_manager)
-        # TODO to manage the non-blocking flows.
+        # TODO: to manage the non-blocking flows.
         producer_task = asyncio.create_task(
             self._run_event_stream(
                 request_context,
@@ -144,7 +146,7 @@ async def on_message_send(
             (
                 result,
                 interrupted,
-            ) = await result_aggregator.consume_and_break_on_interrupt(consumer)
+            ) = await result_aggregator.consume_and_break_on_interrupt(consumer)            
             if not result:
                 raise ServerError(error=InternalError())
         finally:

src/a2a/server/request_handlers/jsonrpc_handler.py

@@ -36,13 +36,15 @@
 )
 from a2a.utils.errors import ServerError
 from a2a.utils.helpers import validate
+from a2a.utils.telemetry import SpanKind, trace_class
 
 
 logger = logging.getLogger(__name__)
 
 
+@trace_class(kind=SpanKind.SERVER)
 class JSONRPCHandler:
-    """A handler that maps the JSONRPC Objects to the request handler and back."""
+    """Maps the JSONRPC Objects to the request handler and back."""
 
     def __init__(
         self,
@@ -53,7 +55,7 @@ def __init__(
 
         Args:
             agent_card: The AgentCard describing the agent's capabilities.
-            request_handler: The handler instance responsible for processing A2A requests.
+            request_handler: The handler instance to process A2A requests.
         """
         self.agent_card = agent_card
         self.request_handler = request_handler

src/a2a/utils/__init__.py

@@ -1,3 +1,6 @@
+# type: ignore
+from opentelemetry.trace import SpanKind as OTelSpanKind
+
 from a2a.utils.artifact import new_text_artifact
 from a2a.utils.helpers import (
     append_artifact_to_task,
@@ -12,7 +15,9 @@
 from a2a.utils.task import new_task
 
 
+SpanKind = OTelSpanKind
 __all__ = [
+    'SpanKind',
     'append_artifact_to_task',
     'build_text_artifact',
     'create_task_obj',

src/a2a/utils/helpers.py

@@ -13,11 +13,13 @@
     TextPart,
 )
 from a2a.utils.errors import ServerError, UnsupportedOperationError
+from a2a.utils.telemetry import trace_function
 
 
 logger = logging.getLogger(__name__)
 
 
+@trace_function()
 def create_task_obj(message_send_params: MessageSendParams) -> Task:
     """Create a new task object from message send params."""
     if not message_send_params.message.contextId:
@@ -31,6 +33,7 @@
     )
 
 
+@trace_function()
 def append_artifact_to_task(task: Task, event: TaskArtifactUpdateEvent) -> None:
     """Helper method for updating Task with new artifact data."""
     if not task.artifacts:
@@ -74,7 +77,7 @@
         # We received a chunk to append, but we don't have an existing artifact.
         # we will ignore this chunk
         logger.warning(
             f'Received append=True for non-existent artifact index {artifact_id} in task {task.id}. Ignoring chunk.'
         )
 
 

src/a2a/utils/telemetry.py

@@ -210,7 +210,11 @@
     return async_wrapper if is_async_func else sync_wrapper
 
 
-def trace_class(include_list: list[str] = None, exclude_list: list[str] = None):
+def trace_class(
+    include_list: list[str] = None,
+    exclude_list: list[str] = None,
+    kind=SpanKind.INTERNAL,
+):
     """A class decorator to automatically trace specified methods of a class.
 
     This decorator iterates over the methods of a class and applies the
@@ -261,10 +265,10 @@
     logger.debug(f'Trace all class {include_list}, {exclude_list}')
     exclude_list = exclude_list or []
 
     def decorator(cls):
         all_methods = {}
         for name, method in inspect.getmembers(cls, inspect.isfunction):
             # Skip Dunders
             if name.startswith('__') and name.endswith('__'):
                 continue
 
@@ -276,9 +280,13 @@
                 continue
 
             all_methods[name] = method
             span_name = f'{cls.__module__}.{cls.__name__}.{name}'
             # Set the decorator on the method.
-            setattr(cls, name, trace_function(span_name=span_name)(method))
+            setattr(
+                cls,
+                name,
+                trace_function(span_name=span_name, kind=kind)(method),
+            )
         return cls
 
     return decorator