Add recording steps and hook to the request handler

Author: mindpower

src/a2a/client/base_client.py

@@ -1,4 +1,5 @@
 from collections.abc import AsyncIterator
+from typing import cast
 
 from a2a.client.client import (
     Client,
@@ -12,6 +13,12 @@
 from a2a.client.middleware import ClientCallInterceptor
 from a2a.client.transports.base import ClientTransport
 from a2a.extensions.base import Extension
+from a2a.extensions.trace import (
+    AgentInvocation,
+    CallTypeEnum,
+    StepAction,
+    TraceExtension,
+)
 from a2a.types import (
     AgentCard,
     GetTaskPushNotificationConfigParams,
@@ -68,9 +75,31 @@ async def send_message(
         Yields:
             An async iterator of `ClientEvent` or a final `Message` response.
         """
+        trace_extension: TraceExtension | None = None
         for extension in self._extensions:
+            if isinstance(extension, TraceExtension):
+                trace_extension = cast(TraceExtension, extension)
             extension.on_client_message(request)
 
+        step = None
+        if trace_extension:
+            trace_id = request.metadata.get('trace', {}).get('trace_id')
+            parent_step_id = request.metadata.get('trace', {}).get(
+                'parent_step_id'
+            )
+            step = trace_extension.start_step(
+                trace_id=trace_id,
+                parent_step_id=parent_step_id,
+                call_type=CallTypeEnum.AGENT,
+                step_action=StepAction(
+                    agent_invocation=AgentInvocation(
+                        agent_url=self._card.url,
+                        agent_name=self._card.name,
+                        requests=request.model_dump(mode='json'),
+                    )
+                ),
+            )
+
         config = MessageSendConfiguration(
             accepted_output_modes=self._config.accepted_output_modes,
             blocking=not self._config.polling,
@@ -82,33 +111,44 @@ async def send_message(
         )
         params = MessageSendParams(message=request, configuration=config)
 
-        if not self._config.streaming or not self._card.capabilities.streaming:
-            response = await self._transport.send_message(
+        try:
+            if (
+                not self._config.streaming
+                or not self._card.capabilities.streaming
+            ):
+                response = await self._transport.send_message(
+                    params, context=context
+                )
+                result = (
+                    (response, None)
+                    if isinstance(response, Task)
+                    else response
+                )
+                await self.consume(result, self._card)
+                yield result
+                return
+
+            tracker = ClientTaskManager()
+            stream = self._transport.send_message_streaming(
                 params, context=context
             )
-            result = (
-                (response, None) if isinstance(response, Task) else response
-            )
-            await self.consume(result, self._card)
-            yield result
-            return
 
-        tracker = ClientTaskManager()
-        stream = self._transport.send_message_streaming(params, context=context)
-
-        first_event = await anext(stream)
-        # The response from a server may be either exactly one Message or a
-        # series of Task updates. Separate out the first message for special
-        # case handling, which allows us to simplify further stream processing.
-        if isinstance(first_event, Message):
-            await self.consume(first_event, self._card)
-            yield first_event
-            return
-
-        yield await self._process_response(tracker, first_event)
-
-        async for event in stream:
-            yield await self._process_response(tracker, event)
+            first_event = await anext(stream)
+            # The response from a server may be either exactly one Message or a
+            # series of Task updates. Separate out the first message for special
+            # case handling, which allows us to simplify further stream processing.
+            if isinstance(first_event, Message):
+                await self.consume(first_event, self._card)
+                yield first_event
+                return
+
+            yield await self._process_response(tracker, first_event)
+
+            async for event in stream:
+                yield await self._process_response(tracker, event)
+        finally:
+            if trace_extension and step:
+                trace_extension.end_step(step.step_id)
 
     async def _process_response(
         self,

src/a2a/extensions/trace.py

@@ -1,6 +1,8 @@
 from __future__ import annotations
 
-from datetime import datetime
+import time
+import uuid
+from datetime import datetime, timezone
 from enum import Enum
 from typing import Any
 
@@ -51,7 +53,7 @@ class Step(A2ABaseModel):
     additional_attributes: dict[str, str] | None = None
     latency: int | None = None
     start_time: datetime
-    end_time: datetime
+    end_time: datetime | None = None
 
 
 class ResponseTrace(A2ABaseModel):
@@ -64,18 +66,81 @@ class ResponseTrace(A2ABaseModel):
 class TraceExtension(Extension):
     """An extension for traceability."""
 
+    def __init__(self, **kwargs: Any) -> None:
+        super().__init__(**kwargs)
+        self.traces: dict[str, ResponseTrace] = {}
+        self._current_steps: dict[str, Step] = {}
+
+    def _generate_id(self, prefix: str) -> str:
+        return f'{prefix}-{uuid.uuid4()}'
+
+    def start_trace(self) -> ResponseTrace:
+        """Starts a new trace."""
+        trace_id = self._generate_id('trace')
+        trace = ResponseTrace(trace_id=trace_id, steps=[])
+        self.traces[trace_id] = trace
+        return trace
+
+    def start_step(
+        self,
+        trace_id: str,
+        parent_step_id: str | None,
+        call_type: CallTypeEnum,
+        step_action: StepAction,
+    ) -> Step:
+        """Starts a new step."""
+        step_id = self._generate_id('step')
+        step = Step(
+            step_id=step_id,
+            trace_id=trace_id,
+            parent_step_id=parent_step_id,
+            call_type=call_type,
+            step_action=step_action,
+            start_time=datetime.now(timezone.utc),
+        )
+        self._current_steps[step_id] = step
+        return step
+
+    def end_step(
+        self,
+        step_id: str,
+        cost: int | None = None,
+        total_tokens: int | None = None,
+        additional_attributes: dict[str, str] | None = None,
+    ) -> None:
+        """Ends a step."""
+        if step_id not in self._current_steps:
+            return
+
+        step = self._current_steps.pop(step_id)
+        step.end_time = datetime.now(timezone.utc)
+        step.latency = int(
+            (step.end_time - step.start_time).total_seconds() * 1000
+        )
+        step.cost = cost
+        step.total_tokens = total_tokens
+        step.additional_attributes = additional_attributes
+
+        if step.trace_id in self.traces:
+            self.traces[step.trace_id].steps.append(step)
+
     def on_client_message(self, message: Any) -> None:
         """Appends trace information to the message."""
-        # This is a placeholder implementation.
+        trace = self.start_trace()
         if message.metadata is None:
             message.metadata = {}
-        message.metadata['trace'] = 'client-trace'
+        message.metadata['trace'] = trace.model_dump(mode='json')
 
     def on_server_message(self, message: Any) -> None:
         """Processes trace information from the message."""
-        # This is a placeholder implementation.
-        if hasattr(message, 'metadata') and 'trace' in message.metadata:
-            print(f"Received trace: {message.metadata['trace']}")
+        if (
+            hasattr(message, 'metadata')
+            and message.metadata is not None
+            and 'trace' in message.metadata
+        ):
+            trace_data = message.metadata['trace']
+            trace = ResponseTrace.model_validate(trace_data)
+            self.traces[trace.trace_id] = trace
 
 
-AgentInvocation.model_rebuild()
+AgentInvocation.model_rebuild()

src/a2a/server/agent_execution/agent_executor.py

@@ -1,4 +1,5 @@
 from abc import ABC, abstractmethod
+from typing import Any
 
 from a2a.server.agent_execution.context import RequestContext
 from a2a.server.events.event_queue import EventQueue
@@ -13,7 +14,10 @@ class AgentExecutor(ABC):
 
     @abstractmethod
     async def execute(
-        self, context: RequestContext, event_queue: EventQueue
+        self,
+        context: RequestContext,
+        event_queue: EventQueue,
+        request_handler: Any,
     ) -> None:
         """Execute the agent's logic for a given request context.
 
@@ -26,6 +30,7 @@ async def execute(
         Args:
             context: The request context containing the message, task ID, etc.
             event_queue: The queue to publish events to.
+            request_handler: The request handler that is executing the agent.
         """
 
     @abstractmethod

src/a2a/server/request_handlers/default_request_handler.py

@@ -27,6 +27,12 @@
     TaskStore,
 )
 from a2a.extensions.base import Extension
+from a2a.extensions.trace import (
+    CallTypeEnum,
+    StepAction,
+    ToolInvocation,
+    TraceExtension,
+)
 from a2a.types import (
     DeleteTaskPushNotificationConfigParams,
     GetTaskPushNotificationConfigParams,
@@ -176,9 +182,40 @@ async def _run_event_stream(
             request: The request context for the agent.
             queue: The event queue for the agent to publish to.
         """
-        await self.agent_executor.execute(request, queue)
+        await self.agent_executor.execute(request, queue, self)
         await queue.close()
 
+    async def handle_tool_call(
+        self,
+        trace_id: str,
+        parent_step_id: str,
+        tool_name: str,
+        parameters: dict[str, Any],
+    ) -> None:
+        """Handles a tool call from the agent executor."""
+        trace_extension: TraceExtension | None = None
+        for extension in self._extensions:
+            if isinstance(extension, TraceExtension):
+                trace_extension = cast(TraceExtension, extension)
+
+        if not trace_extension:
+            return
+
+        step = trace_extension.start_step(
+            trace_id=trace_id,
+            parent_step_id=parent_step_id,
+            call_type=CallTypeEnum.TOOL,
+            step_action=StepAction(
+                tool_invocation=ToolInvocation(
+                    tool_name=tool_name,
+                    parameters=parameters,
+                )
+            ),
+        )
+        # In a real implementation, you would execute the tool here.
+        # For this example, we'll just end the step immediately.
+        trace_extension.end_step(step.step_id)
+
     async def _setup_message_execution(
         self,
         params: MessageSendParams,

tests/extensions/debug_trace.py

@@ -0,0 +1,51 @@
+#!/usr/bin/env python3
+import sys
+from unittest.mock import Mock
+
+from a2a.client.base_client import BaseClient
+from a2a.extensions.trace import TraceExtension
+from a2a.types import Message, TextPart, Role, Part
+
+def debug_trace():
+    print("Starting trace debug...")
+    
+    # Create the extension
+    trace_extension = TraceExtension()
+    
+    # Create a trace directly to see its structure
+    trace = trace_extension.start_trace()
+    print(f"Direct trace object: {trace}")
+    print(f"Direct trace dict: {trace.model_dump(mode='json')}")
+    
+    # Create a message
+    message = Message(
+        message_id='test_message',
+        role=Role.user,
+        parts=[Part(TextPart(text='Hello, world!'))],
+    )
+    
+    print(f"Initial message metadata: {message.metadata}")
+    
+    # Call the extension method
+    trace_extension.on_client_message(message)
+    
+    print(f"After extension metadata: {message.metadata}")
+    
+    if message.metadata and 'trace' in message.metadata:
+        trace_data = message.metadata['trace']
+        print(f"Trace data type: {type(trace_data)}")
+        print(f"Trace data: {trace_data}")
+        
+        if isinstance(trace_data, dict):
+            print(f"Trace data keys: {list(trace_data.keys())}")
+            if 'trace_id' in trace_data:
+                print(f"Found trace_id: {trace_data['trace_id']}")
+            else:
+                print("trace_id not found in trace data")
+        else:
+            print("Trace data is not a dict")
+    else:
+        print("No trace data found in metadata")
+
+if __name__ == "__main__":
+    debug_trace()

tests/extensions/simple_trace_test.py

@@ -0,0 +1,22 @@
+#!/usr/bin/env python3
+
+# Simple test to check ResponseTrace serialization
+from a2a.extensions.trace import TraceExtension, ResponseTrace
+
+# Create extension and trace
+ext = TraceExtension()
+trace = ext.start_trace()
+
+print("Trace object:", trace)
+print("Trace type:", type(trace))
+print("Trace fields:", trace.__dict__)
+print("Model dump:", trace.model_dump(mode='json'))
+
+# Test creating trace data like in the extension
+if True:  # message.metadata is None
+    metadata = {}
+metadata['trace'] = trace.model_dump(mode='json')
+
+print("Metadata:", metadata)
+print("Trace in metadata:", metadata['trace'])
+print("Keys in trace:", metadata['trace'].keys() if isinstance(metadata['trace'], dict) else "not a dict")