From bf3f425f5ca36b1bd30c8bcd9e42dda35366d6c2 Mon Sep 17 00:00:00 2001 From: Anik Bhattacharjee Date: Fri, 21 Aug 2026 16:25:05 -0400 Subject: [PATCH] LCORE-2984: Add OpenTelemetry spans to A2A endpoint - Adds `a2a.dispatch` span to `_handle_a2a_jsonrpc` covering the full JSON-RPC dispatch lifecycle (both streaming and non-streaming paths), with `a2a.rpc.method`, `a2a.request.id`, and anonymized `user.id` attributes plus `a2a.dispatch.start`/`a2a.dispatch.end` events. - Adds `a2a.execute` span to `_process_task_streaming` capturing `session.id`, anonymized `input`/`output`, `llm.model.id`, `llm.provider.id`, `llm.usage.input_tokens`, `llm.usage.output_tokens`, `tool.calls.count`, `tool.calls.names`, with `llm.inference.completed` and `tool.execution.completed` events. - Adds `A2A_RPC_METHOD`, `A2A_REQUEST_ID` to `SpanAttributes` and `A2A_DISPATCH_START`, `A2A_DISPATCH_END` to `SpanEvents` enums in `otel_tracing.py`. - Tool call tracking is wired through `_convert_stream_to_events` via a mutable list that collects tool names from `FunctionToolCallEvent` and `NativeToolCallPart` stream events. Signed-off-by: Anik Bhattacharjee --- src/app/endpoints/a2a.py | 670 ++++++++++++++++----------- src/utils/otel_tracing.py | 4 + tests/unit/app/endpoints/test_a2a.py | 529 +++++++++++++++++++++ 3 files changed, 924 insertions(+), 279 deletions(-) diff --git a/src/app/endpoints/a2a.py b/src/app/endpoints/a2a.py index 60f83cba4..09ced0482 100644 --- a/src/app/endpoints/a2a.py +++ b/src/app/endpoints/a2a.py @@ -34,6 +34,7 @@ from a2a.utils import new_agent_text_message, new_task from fastapi import APIRouter, Depends, HTTPException, Request, status from ogx_client import APIConnectionError, APIStatusError +from opentelemetry import trace from pydantic_ai import AgentRunResultEvent from pydantic_ai.exceptions import AgentRunError from pydantic_ai.messages import ( @@ -63,12 +64,21 @@ from utils.agents.error_handler import map_agent_inference_error from utils.conversation_compaction import apply_compaction_blocking from utils.mcp_headers import McpHeaders, mcp_headers_dependency +from utils.otel_tracing import ( + SpanAttributes, + SpanEvents, + add_span_event, + anonymize_value, + set_span_attributes, +) from utils.pydantic_ai_helpers import build_agent +from utils.query import extract_provider_and_model_from_model_id from utils.responses import prepare_responses_params from utils.suid import normalize_conversation_id from version import __version__ logger = get_logger(__name__) +tracer = trace.get_tracer(__name__) router = APIRouter(tags=["a2a"]) auth_dependency = get_auth_dependency() @@ -136,6 +146,61 @@ def _build_a2a_parts_from_agent_result( return [Part(root=A2ATextPart(text=final_text))] +def _record_model_span(span: trace.Span, model_id: str) -> None: + """Set LLM model and provider attributes on a span. + + Parameters: + span: The active OpenTelemetry span. + model_id: Full model identifier in "provider/model" format. + """ + provider_id, _ = extract_provider_and_model_from_model_id(model_id) + set_span_attributes( + span, + { + SpanAttributes.LLM_MODEL_ID: model_id, + SpanAttributes.LLM_PROVIDER_ID: provider_id, + }, + ) + + +def _record_execution_span( + span: trace.Span, + tool_call_names: list[str], + run_result: Optional[AgentRunResult[str]], +) -> None: + """Record tool-call metrics, token usage, and output on an a2a.execute span. + + Parameters: + span: The active OpenTelemetry span. + tool_call_names: Tool names collected during streaming. + run_result: Completed agent run result, or None. + """ + if tool_call_names: + set_span_attributes( + span, + { + SpanAttributes.TOOL_CALLS_COUNT: len(tool_call_names), + SpanAttributes.TOOL_CALLS_NAMES: ",".join(sorted(set(tool_call_names))), + }, + ) + add_span_event(span, SpanEvents.TOOL_EXECUTION_COMPLETED) + + if run_result is not None: + usage = run_result.usage + set_span_attributes( + span, + { + SpanAttributes.LLM_USAGE_INPUT_TOKENS: usage.input_tokens, + SpanAttributes.LLM_USAGE_OUTPUT_TOKENS: usage.output_tokens, + }, + ) + add_span_event(span, SpanEvents.LLM_INFERENCE_COMPLETED) + + output_text = run_result.response.text + if output_text: + span.set_attribute(SpanAttributes.OUTPUT, anonymize_value(output_text)) + + class TaskResultAggregator: """Aggregates the task status updates and provides the final task state.""" @@ -223,6 +288,7 @@ def __init__( self.mcp_headers: McpHeaders = mcp_headers or {} self.request_headers: Optional[Mapping[str, str]] = request_headers self._run_result: Optional[AgentRunResult[str]] = None + self._tool_call_names: list[str] = [] async def execute( self, @@ -292,179 +358,195 @@ async def _process_task_streaming( # pylint: disable=too-many-locals if not task_id or not context_id: raise ValueError("Task ID and Context ID are required") - # Extract user input using SDK utility - user_input = context.get_user_input() - if not user_input: - await task_updater.update_status( - TaskState.input_required, - message=new_agent_text_message( - "No input received. Please provide your input.", - context_id=context_id, - task_id=task_id, - ), - final=True, - ) - return - - preview = user_input[:200] + ("..." if len(user_input) > 200 else "") - logger.info("Processing A2A request: %s", preview) - - # Extract routing metadata from A2A message context. - # Supported metadata fields (see docs/a2a_protocol.md for details): - # - model: LLM model to use (e.g., "gpt-4", "llama3.1") - # - provider: LLM provider to use (e.g., "openai", "watsonx") - # - vector_store_ids: list of vector store IDs for RAG queries - metadata = context.message.metadata if context.message else {} - model = metadata.get("model") if metadata else None - provider = metadata.get("provider") if metadata else None - vector_store_ids = metadata.get("vector_store_ids") if metadata else None - - # Resolve conversation_id from A2A contextId for multi-turn - a2a_context_id = context_id - context_store = await _get_context_store() - conversation_id = await context_store.get(a2a_context_id) - logger.info( - "A2A contextId %s maps to conversation_id %s", - a2a_context_id, - conversation_id, - ) - - # Build internal query request (conversation_id may be None for first turn) - query_request = QueryRequest( - query=user_input, - conversation_id=conversation_id, - model=model, - provider=provider, - system_prompt=None, - attachments=None, - no_tools=False, - generate_topic_summary=True, - media_type=None, - vector_store_ids=vector_store_ids, - shield_ids=None, - solr=None, - ) + with tracer.start_as_current_span("a2a.execute") as span: + span.set_attribute(SpanAttributes.SESSION_ID, context_id) - # Get LLM client and select model - client = AsyncOgxClientHolder().get_client() - try: - responses_params = await prepare_responses_params( - client, - query_request, - None, - self.auth_token, - self.mcp_headers, - stream=True, - store=True, - request_headers=self.request_headers, - ) - # Compact the conversation if it is approaching the context window - # limit. A2A is not a browser SSE stream, so no progress event is - # emitted; the blocking variant summarizes inline before the call. - # No conversation cache is passed: the A2A executor has no resolved - # user_id for the (user_id, conversation_id) cache key, so A2A runs - # in marker-only mode (additive summaries, no persisted fold). - compaction = await apply_compaction_blocking( - client, - responses_params, - configuration.inference, - configuration.compaction, + # Extract user input using SDK utility + user_input = context.get_user_input() + if not user_input: + await task_updater.update_status( + TaskState.input_required, + message=new_agent_text_message( + "No input received. Please provide your input.", + context_id=context_id, + task_id=task_id, + ), + final=True, + ) + return + + span.set_attribute(SpanAttributes.INPUT, anonymize_value(user_input)) + preview = user_input[:200] + ("..." if len(user_input) > 200 else "") + logger.info("Processing A2A request: %s", preview) + + # Extract routing metadata from A2A message context. + # Supported metadata fields (see docs/a2a_protocol.md for details): + # - model: LLM model to use (e.g., "gpt-4", "llama3.1") + # - provider: LLM provider to use (e.g., "openai", "watsonx") + # - vector_store_ids: list of vector store IDs for RAG queries + metadata = context.message.metadata if context.message else {} + + # Resolve conversation_id from A2A contextId for multi-turn + context_store = await _get_context_store() + conversation_id = await context_store.get(context_id) + logger.info( + "A2A contextId %s maps to conversation_id %s", + context_id, + conversation_id, ) - responses_params = compaction.params - agent = build_agent( - client, - responses_params, - configuration, - shields=query_request.shield_ids, - ) - except (AgentRunError, APIStatusError, APIConnectionError, RuntimeError) as e: - error_response = map_agent_inference_error(e, query_request.model or "") - logger.error("Error preparing A2A agent: %s", str(e), exc_info=True) - await task_updater.update_status( - TaskState.failed, - message=new_agent_text_message( - error_response.detail.response, - context_id=context_id, - task_id=task_id, + # Build internal query request (conversation_id may be None for first turn) + query_request = QueryRequest( + query=user_input, + conversation_id=conversation_id, + model=metadata.get("model") if metadata else None, + provider=metadata.get("provider") if metadata else None, + system_prompt=None, + attachments=None, + no_tools=False, + generate_topic_summary=True, + media_type=None, + vector_store_ids=( + metadata.get("vector_store_ids") if metadata else None ), - final=True, + shield_ids=None, + solr=None, ) - return - # Persist conversation_id for next turn in same A2A context - conversation_id = conversation_id or normalize_conversation_id( - responses_params.conversation - ) - if conversation_id: - await context_store.set(a2a_context_id, conversation_id) - logger.info( - "Persisted conversation_id %s for A2A contextId %s", - conversation_id, - a2a_context_id, - ) - - # Initialize result aggregator - aggregator = TaskResultAggregator() - event_queue = task_updater.event_queue + # Get LLM client and select model + client = AsyncOgxClientHolder().get_client() + try: + responses_params = await prepare_responses_params( + client, + query_request, + None, + self.auth_token, + self.mcp_headers, + stream=True, + store=True, + request_headers=self.request_headers, + ) + # Compact the conversation if it is approaching the context window + # limit. A2A is not a browser SSE stream, so no progress event is + # emitted; the blocking variant summarizes inline before the call. + # No conversation cache is passed: the A2A executor has no resolved + # user_id for the (user_id, conversation_id) cache key, so A2A runs + # in marker-only mode (additive summaries, no persisted fold). + compaction = await apply_compaction_blocking( + client, + responses_params, + configuration.inference, + configuration.compaction, + ) + responses_params = compaction.params + + _record_model_span(span, responses_params.model) + agent = build_agent( + client, + responses_params, + configuration, + shields=query_request.shield_ids, + ) + except ( + AgentRunError, + APIStatusError, + APIConnectionError, + RuntimeError, + ) as e: + error_response = map_agent_inference_error(e, query_request.model or "") + logger.error("Error preparing A2A agent: %s", str(e), exc_info=True) + await task_updater.update_status( + TaskState.failed, + message=new_agent_text_message( + error_response.detail.response, + context_id=context_id, + task_id=task_id, + ), + final=True, + ) + return - # Emit working status with metadata before processing stream - await event_queue.enqueue_event( - TaskStatusUpdateEvent( - task_id=task_id, - status=TaskStatus( - state=TaskState.working, - timestamp=datetime.now(UTC).isoformat(), - ), - context_id=context_id, - final=False, - metadata={ - "model": responses_params.model, - "conversation_id": conversation_id, - }, + # Persist conversation_id for next turn in same A2A context + conversation_id = conversation_id or normalize_conversation_id( + responses_params.conversation ) - ) + if conversation_id: + await context_store.set(context_id, conversation_id) + logger.info( + "Persisted conversation_id %s for A2A contextId %s", + conversation_id, + context_id, + ) - # Run the pydantic-ai agent and convert stream events to A2A events. - prompt = user_input - try: - async for a2a_event in self._convert_stream_to_events( - agent, - prompt, - task_id, - context_id, - conversation_id=conversation_id, - ): - aggregator.process_event(a2a_event) - await event_queue.enqueue_event(a2a_event) - except (AgentRunError, APIStatusError, APIConnectionError, RuntimeError) as e: - error_response = map_agent_inference_error(e, responses_params.model) - logger.error("Error during A2A agent run: %s", str(e), exc_info=True) - await task_updater.update_status( - TaskState.failed, - message=new_agent_text_message( - error_response.detail.response, - context_id=context_id, + # Initialize result aggregator + aggregator = TaskResultAggregator() + event_queue = task_updater.event_queue + + # Emit working status with metadata before processing stream + await event_queue.enqueue_event( + TaskStatusUpdateEvent( task_id=task_id, - ), - final=True, + status=TaskStatus( + state=TaskState.working, + timestamp=datetime.now(UTC).isoformat(), + ), + context_id=context_id, + final=False, + metadata={ + "model": responses_params.model, + "conversation_id": conversation_id, + }, + ) ) - return - # Publish the final task result event - if aggregator.task_state == TaskState.working: - await task_updater.update_status( - TaskState.completed, - timestamp=datetime.now(UTC).isoformat(), - final=True, - ) - else: - await task_updater.update_status( - aggregator.task_state, - message=aggregator.task_status_message, - timestamp=datetime.now(UTC).isoformat(), - final=True, - ) + # Run the pydantic-ai agent and convert stream events to A2A events. + prompt = user_input + self._tool_call_names = [] + try: + async for a2a_event in self._convert_stream_to_events( + agent, + prompt, + task_id, + context_id, + conversation_id=conversation_id, + ): + aggregator.process_event(a2a_event) + await event_queue.enqueue_event(a2a_event) + except ( + AgentRunError, + APIStatusError, + APIConnectionError, + RuntimeError, + ) as e: + error_response = map_agent_inference_error(e, responses_params.model) + logger.error("Error during A2A agent run: %s", str(e), exc_info=True) + await task_updater.update_status( + TaskState.failed, + message=new_agent_text_message( + error_response.detail.response, + context_id=context_id, + task_id=task_id, + ), + final=True, + ) + return + + _record_execution_span(span, self._tool_call_names, self._run_result) + + # Publish the final task result event + if aggregator.task_state == TaskState.working: + await task_updater.update_status( + TaskState.completed, + timestamp=datetime.now(UTC).isoformat(), + final=True, + ) + else: + await task_updater.update_status( + aggregator.task_state, + message=aggregator.task_status_message, + timestamp=datetime.now(UTC).isoformat(), + final=True, + ) async def _convert_stream_to_events( self, @@ -499,6 +581,12 @@ async def _convert_stream_to_events( run_result = event.result self._run_result = run_result continue + if isinstance(event, FunctionToolCallEvent): + self._tool_call_names.append(event.part.tool_name) + elif isinstance(event, PartEndEvent) and isinstance( + event.part, NativeToolCallPart + ): + self._tool_call_names.append(event.part.tool_name) a2a_event = self._dispatch_agent_event( event, task_id, context_id, text_parts, artifact_id ) @@ -928,6 +1016,8 @@ async def _handle_a2a_jsonrpc( # pylint: disable=too-many-locals,too-many-state # Detect if this is a streaming request by checking the JSON-RPC method is_streaming_request = False + rpc_method = "" + rpc_request_id = "" body = b"" try: # Read and parse the request body to check the method @@ -937,11 +1027,12 @@ async def _handle_a2a_jsonrpc( # pylint: disable=too-many-locals,too-many-state try: rpc_request = json.loads(body) # Check if the method is message/stream - method = rpc_request.get("method", "") - is_streaming_request = method == "message/stream" + rpc_method = rpc_request.get("method", "") + rpc_request_id = str(rpc_request.get("id", "")) + is_streaming_request = rpc_method == "message/stream" logger.info( "A2A request method: %s, streaming: %s", - method, + rpc_method, is_streaming_request, ) except (json.JSONDecodeError, AttributeError) as e: @@ -951,124 +1042,145 @@ async def _handle_a2a_jsonrpc( # pylint: disable=too-many-locals,too-many-state except Exception as e: # pylint: disable=broad-except logger.error("Error detecting streaming request: %s", str(e)) - # Setup scope for A2A app - scope = dict(request.scope) - scope["path"] = "/" # A2A app expects root path - - # We need to re-provide the body since we already read it - body_sent = False - - async def receive() -> MutableMapping[str, Any]: - nonlocal body_sent - if not body_sent: - body_sent = True - return {"type": "http.request", "body": body, "more_body": False} - - # After sending body once, delegate to original receive - # This prevents infinite loops - the original receive() will block/disconnect properly - return await request.receive() - - if is_streaming_request: - # Streaming mode: Forward chunks to client as they arrive - logger.info("Handling A2A streaming request") - - # Create queue for passing chunks from ASGI app to response generator - chunk_queue: asyncio.Queue[Optional[bytes]] = asyncio.Queue() - - async def streaming_send(message: dict[str, Any]) -> None: - """Send callback that queues chunks for streaming.""" - if message["type"] == "http.response.body": - body_chunk = message.get("body", b"") - if body_chunk: - await chunk_queue.put(body_chunk) - # Signal end of stream if no more body - if not message.get("more_body", False): - logger.debug("Streaming: End of stream signaled") - await chunk_queue.put(None) - - # Run the A2A app in a background task - async def run_a2a_app() -> None: - """Run A2A app and handle any errors.""" - try: - logger.debug("Streaming: Starting A2A app execution") - await a2a_app(scope, receive, streaming_send) - logger.debug("Streaming: A2A app execution completed") - except Exception as exc: # pylint: disable=broad-except - logger.error( - "Error in A2A app during streaming: %s", str(exc), exc_info=True - ) - await chunk_queue.put(None) # Signal end even on error + with tracer.start_as_current_span("a2a.dispatch") as span: + set_span_attributes( + span, + { + SpanAttributes.A2A_RPC_METHOD: rpc_method, + SpanAttributes.A2A_REQUEST_ID: ( + anonymize_value(rpc_request_id) if rpc_request_id else "" + ), + SpanAttributes.USER_ID: anonymize_value(auth[0]) if auth[0] else "", + }, + ) + add_span_event(span, SpanEvents.A2A_DISPATCH_START) + + # Setup scope for A2A app + scope = dict(request.scope) + scope["path"] = "/" # A2A app expects root path + + # We need to re-provide the body since we already read it + body_sent = False + + async def receive() -> MutableMapping[str, Any]: + nonlocal body_sent + if not body_sent: + body_sent = True + return {"type": "http.request", "body": body, "more_body": False} + + # After sending body once, delegate to original receive + # This prevents infinite loops - the original receive() will block/disconnect properly + return await request.receive() + + if is_streaming_request: + # Streaming mode: Forward chunks to client as they arrive + logger.info("Handling A2A streaming request") + + # Create queue for passing chunks from ASGI app to response generator + chunk_queue: asyncio.Queue[Optional[bytes]] = asyncio.Queue() + + async def streaming_send(message: dict[str, Any]) -> None: + """Send callback that queues chunks for streaming.""" + if message["type"] == "http.response.body": + body_chunk = message.get("body", b"") + if body_chunk: + await chunk_queue.put(body_chunk) + # Signal end of stream if no more body + if not message.get("more_body", False): + logger.debug("Streaming: End of stream signaled") + await chunk_queue.put(None) + + # Run the A2A app in a background task + async def run_a2a_app() -> None: + """Run A2A app and handle any errors.""" + try: + logger.debug("Streaming: Starting A2A app execution") + await a2a_app(scope, receive, streaming_send) + logger.debug("Streaming: A2A app execution completed") + except Exception as exc: # pylint: disable=broad-except + logger.error( + "Error in A2A app during streaming: %s", + str(exc), + exc_info=True, + ) + await chunk_queue.put(None) # Signal end even on error + + # Start the A2A app task + app_task = asyncio.create_task(run_a2a_app()) + + async def response_generator() -> AsyncIterator[bytes]: + """Generate chunks from the queue for streaming response.""" + chunk_count = 0 + try: + while True: + # Get chunk from queue with timeout to prevent hanging + try: + chunk = await asyncio.wait_for( + chunk_queue.get(), timeout=300.0 + ) + except TimeoutError: + logger.error("Timeout waiting for chunk from A2A app") + break + + if chunk is None: + # End of stream + logger.debug( + "Streaming: Stream ended after %d chunks", chunk_count + ) + break + chunk_count += 1 + logger.debug("Chunk sent to A2A client: %s", str(chunk)) + yield chunk + finally: + # Ensure the app task is cleaned up + if not app_task.done(): + app_task.cancel() + try: + await app_task + except asyncio.CancelledError: + pass + + # Return streaming response immediately + # The status code and headers will be determined by the first chunk + # We can't wait for the response to start because that would cause a deadlock: + # the ASGI app won't send data until the client starts consuming + logger.debug("Streaming: Returning StreamingResponse") + + add_span_event(span, SpanEvents.A2A_DISPATCH_END) + + # Return streaming response with SSE content type for A2A protocol + return StreamingResponse( + response_generator(), + media_type=MEDIA_TYPE_EVENT_STREAM, + ) - # Start the A2A app task - app_task = asyncio.create_task(run_a2a_app()) + # Non-streaming mode: Buffer entire response + logger.info("Handling A2A non-streaming request") - async def response_generator() -> AsyncIterator[bytes]: - """Generate chunks from the queue for streaming response.""" - chunk_count = 0 - try: - while True: - # Get chunk from queue with timeout to prevent hanging - try: - chunk = await asyncio.wait_for(chunk_queue.get(), timeout=300.0) - except TimeoutError: - logger.error("Timeout waiting for chunk from A2A app") - break - - if chunk is None: - # End of stream - logger.debug( - "Streaming: Stream ended after %d chunks", chunk_count - ) - break - chunk_count += 1 - logger.debug("Chunk sent to A2A client: %s", str(chunk)) - yield chunk - finally: - # Ensure the app task is cleaned up - if not app_task.done(): - app_task.cancel() - try: - await app_task - except asyncio.CancelledError: - pass - - # Return streaming response immediately - # The status code and headers will be determined by the first chunk - # We can't wait for the response to start because that would cause a deadlock: - # the ASGI app won't send data until the client starts consuming - logger.debug("Streaming: Returning StreamingResponse") - - # Return streaming response with SSE content type for A2A protocol - return StreamingResponse( - response_generator(), - media_type=MEDIA_TYPE_EVENT_STREAM, - ) + response_started = False + response_body = [] + status_code = 200 + headers = [] - # Non-streaming mode: Buffer entire response - logger.info("Handling A2A non-streaming request") - - response_started = False - response_body = [] - status_code = 200 - headers = [] - - async def buffering_send(message: dict[str, Any]) -> None: - nonlocal response_started, status_code, headers - if message["type"] == "http.response.start": - response_started = True - status_code = message["status"] - headers = message.get("headers", []) - elif message["type"] == "http.response.body": - response_body.append(message.get("body", b"")) - - await a2a_app(scope, receive, buffering_send) - - # Return the response from A2A app - return Response( - content=b"".join(response_body), - status_code=status_code, - headers=dict((k.decode(), v.decode()) for k, v in headers), - ) + async def buffering_send(message: dict[str, Any]) -> None: + nonlocal response_started, status_code, headers + if message["type"] == "http.response.start": + response_started = True + status_code = message["status"] + headers = message.get("headers", []) + elif message["type"] == "http.response.body": + response_body.append(message.get("body", b"")) + + await a2a_app(scope, receive, buffering_send) + + add_span_event(span, SpanEvents.A2A_DISPATCH_END) + + # Return the response from A2A app + return Response( + content=b"".join(response_body), + status_code=status_code, + headers=dict((k.decode(), v.decode()) for k, v in headers), + ) @router.get("/a2a/health") diff --git a/src/utils/otel_tracing.py b/src/utils/otel_tracing.py index 1fdf3ca54..65c5f6fb0 100644 --- a/src/utils/otel_tracing.py +++ b/src/utils/otel_tracing.py @@ -43,6 +43,8 @@ class SpanAttributes(StrEnum): SKILL_ACTIVATIONS = "skill.activations" RLS_TEMPLATE_OK = "rls.template.ok" TOPIC_SUMMARY_SUCCESS = "topic.summary.success" + A2A_RPC_METHOD = "a2a.rpc.method" + A2A_REQUEST_ID = "a2a.request.id" class SpanEvents(StrEnum): @@ -61,6 +63,8 @@ class SpanEvents(StrEnum): TURN_PERSISTED = "turn.persisted" TOPIC_SUMMARY_TASK_STARTED = "topic.summary.task.started" TOPIC_SUMMARY_TASK_FINISHED = "topic.summary.task.finished" + A2A_DISPATCH_START = "a2a.dispatch.start" + A2A_DISPATCH_END = "a2a.dispatch.end" def anonymize_value(value: str, max_length: int = 50) -> str: diff --git a/tests/unit/app/endpoints/test_a2a.py b/tests/unit/app/endpoints/test_a2a.py index 308e185f8..aed44df2f 100644 --- a/tests/unit/app/endpoints/test_a2a.py +++ b/tests/unit/app/endpoints/test_a2a.py @@ -24,6 +24,9 @@ from fastapi import HTTPException, Request from ogx_client import APIConnectionError from ogx_client.types import ListModelsResponse +from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( + InMemorySpanExporter, +) from pydantic_ai import AgentRunResultEvent from pydantic_ai.exceptions import AgentRunError from pydantic_ai.messages import ( @@ -44,6 +47,7 @@ _build_a2a_parts_from_agent_result, _get_context_store, _get_task_store, + _handle_a2a_jsonrpc, a2a_health_check, get_agent_card, get_lightspeed_agent_card, @@ -1223,3 +1227,528 @@ async def test_get_agent_card_endpoint( assert isinstance(result, AgentCard) assert result.name == "Test Agent" assert result.url == "http://localhost:8080/a2a" + + +# ----------------------------- +# Tests for A2A OTEL Spans +# ----------------------------- +class TestA2AOtelSpans: + """Tests for OpenTelemetry span instrumentation on A2A endpoints.""" + + @pytest.mark.asyncio + async def test_execute_span_success_attributes( # pylint: disable=too-many-locals,too-many-statements + self, + mocker: MockerFixture, + setup_configuration: AppConfig, # pylint: disable=unused-argument + otel: tuple[Any, InMemorySpanExporter], + ) -> None: + """Test that a2a.execute span records model, token usage, and output.""" + tracer, exporter = otel + mocker.patch("app.endpoints.a2a.tracer", tracer) + + executor = A2AAgentExecutor(auth_token="test-token") + + mock_message = mocker.MagicMock() + mock_message.role = "user" + mock_message.parts = [Part(root=TextPart(text="Hello"))] + mock_message.metadata = {} + + context = mocker.MagicMock(spec=RequestContext) + context.task_id = "task-123" + context.context_id = "ctx-456" + context.message = mock_message + context.get_user_input.return_value = "Hello A2A" + + event_queue = mocker.AsyncMock(spec=EventQueue) + task_updater = mocker.MagicMock() + task_updater.update_status = mocker.AsyncMock() + task_updater.event_queue = event_queue + + mock_context_store = mocker.AsyncMock() + mock_context_store.get.return_value = None + mocker.patch( + "app.endpoints.a2a._get_context_store", return_value=mock_context_store + ) + + mock_client = mocker.AsyncMock() + mock_client.models.list = mocker.AsyncMock( + return_value=ListModelsResponse.model_construct(data=[mocker.MagicMock()]) + ) + mocker.patch( + "app.endpoints.a2a.AsyncOgxClientHolder" + ).return_value.get_client.return_value = mock_client + + mock_responses_params = mocker.Mock() + mock_responses_params.model = "watsonx/granite-3.1" + mock_responses_params.conversation = "conv_x" + mocker.patch( + "app.endpoints.a2a.prepare_responses_params", + new=mocker.AsyncMock(return_value=mock_responses_params), + ) + + compaction_result = mocker.Mock() + compaction_result.params = mock_responses_params + mocker.patch( + "app.endpoints.a2a.apply_compaction_blocking", + new=mocker.AsyncMock(return_value=compaction_result), + ) + + mock_run_result = mocker.MagicMock() + mock_run_result.response.text = "A2A response text" + mock_usage = mocker.MagicMock() + mock_usage.input_tokens = 42 + mock_usage.output_tokens = 18 + mock_run_result.usage = mock_usage + + result_event = mocker.MagicMock(spec=AgentRunResultEvent) + result_event.result = mock_run_result + + async def _event_stream() -> Any: + yield result_event + + mock_stream_ctx = mocker.AsyncMock() + mock_stream_ctx.__aenter__ = mocker.AsyncMock(return_value=_event_stream()) + mock_stream_ctx.__aexit__ = mocker.AsyncMock(return_value=False) + mock_agent = mocker.MagicMock() + mock_agent.run_stream_events.return_value = mock_stream_ctx + mocker.patch("app.endpoints.a2a.build_agent", return_value=mock_agent) + + await executor._process_task_streaming( + context, task_updater, context.task_id, context.context_id + ) + + spans = exporter.get_finished_spans() + execute_spans = [s for s in spans if s.name == "a2a.execute"] + assert len(execute_spans) == 1 + span = execute_spans[0] + attrs = dict(span.attributes or {}) + + assert attrs["session.id"] == "ctx-456" + assert attrs["llm.model.id"] == "watsonx/granite-3.1" + assert attrs["llm.provider.id"] == "watsonx" + assert attrs["llm.usage.input_tokens"] == 42 + assert attrs["llm.usage.output_tokens"] == 18 + assert "request.input" in attrs + assert "response.output" in attrs + + @pytest.mark.asyncio + async def test_execute_span_tool_calls( # pylint: disable=too-many-locals,too-many-statements + self, + mocker: MockerFixture, + setup_configuration: AppConfig, # pylint: disable=unused-argument + otel: tuple[Any, InMemorySpanExporter], + ) -> None: + """Test that tool calls are tracked on the a2a.execute span.""" + tracer, exporter = otel + mocker.patch("app.endpoints.a2a.tracer", tracer) + + executor = A2AAgentExecutor(auth_token="test-token") + + mock_message = mocker.MagicMock() + mock_message.role = "user" + mock_message.parts = [Part(root=TextPart(text="Hello"))] + mock_message.metadata = {} + + context = mocker.MagicMock(spec=RequestContext) + context.task_id = "task-123" + context.context_id = "ctx-456" + context.message = mock_message + context.get_user_input.return_value = "Use tools" + + event_queue = mocker.AsyncMock(spec=EventQueue) + task_updater = mocker.MagicMock() + task_updater.update_status = mocker.AsyncMock() + task_updater.event_queue = event_queue + + mock_context_store = mocker.AsyncMock() + mock_context_store.get.return_value = None + mocker.patch( + "app.endpoints.a2a._get_context_store", return_value=mock_context_store + ) + + mock_client = mocker.AsyncMock() + mock_client.models.list = mocker.AsyncMock( + return_value=ListModelsResponse.model_construct(data=[mocker.MagicMock()]) + ) + mocker.patch( + "app.endpoints.a2a.AsyncOgxClientHolder" + ).return_value.get_client.return_value = mock_client + + mock_responses_params = mocker.Mock() + mock_responses_params.model = "openai/gpt-4" + mock_responses_params.conversation = "conv_x" + mocker.patch( + "app.endpoints.a2a.prepare_responses_params", + new=mocker.AsyncMock(return_value=mock_responses_params), + ) + + compaction_result = mocker.Mock() + compaction_result.params = mock_responses_params + mocker.patch( + "app.endpoints.a2a.apply_compaction_blocking", + new=mocker.AsyncMock(return_value=compaction_result), + ) + + # Stream events: two FunctionToolCallEvents + result + tool_event_1 = mocker.MagicMock(spec=FunctionToolCallEvent) + tool_event_1.part = mocker.MagicMock() + tool_event_1.part.tool_name = "search_docs" + tool_event_1.part.tool_call_id = "tc1" + + tool_event_2 = mocker.MagicMock(spec=FunctionToolCallEvent) + tool_event_2.part = mocker.MagicMock() + tool_event_2.part.tool_name = "get_weather" + tool_event_2.part.tool_call_id = "tc2" + + mock_run_result = mocker.MagicMock() + mock_run_result.response.text = "Tool result" + mock_usage = mocker.MagicMock() + mock_usage.input_tokens = 100 + mock_usage.output_tokens = 50 + mock_run_result.usage = mock_usage + + result_event = mocker.MagicMock(spec=AgentRunResultEvent) + result_event.result = mock_run_result + + async def _event_stream() -> Any: + yield tool_event_1 + yield tool_event_2 + yield result_event + + mock_stream_ctx = mocker.AsyncMock() + mock_stream_ctx.__aenter__ = mocker.AsyncMock(return_value=_event_stream()) + mock_stream_ctx.__aexit__ = mocker.AsyncMock(return_value=False) + mock_agent = mocker.MagicMock() + mock_agent.run_stream_events.return_value = mock_stream_ctx + mocker.patch("app.endpoints.a2a.build_agent", return_value=mock_agent) + + await executor._process_task_streaming( + context, task_updater, context.task_id, context.context_id + ) + + spans = exporter.get_finished_spans() + execute_spans = [s for s in spans if s.name == "a2a.execute"] + assert len(execute_spans) == 1 + span = execute_spans[0] + attrs = dict(span.attributes or {}) + + assert attrs["tool.calls.count"] == 2 + assert "get_weather" in attrs["tool.calls.names"] + assert "search_docs" in attrs["tool.calls.names"] + + events = span.events + event_names = [e.name for e in events] + assert "tool.execution.completed" in event_names + assert "llm.inference.completed" in event_names + + @pytest.mark.asyncio + async def test_execute_span_inference_completed_event( # pylint: disable=too-many-locals,too-many-statements + self, + mocker: MockerFixture, + setup_configuration: AppConfig, # pylint: disable=unused-argument + otel: tuple[Any, InMemorySpanExporter], + ) -> None: + """Test that llm.inference.completed event is emitted when result is available.""" + tracer, exporter = otel + mocker.patch("app.endpoints.a2a.tracer", tracer) + + executor = A2AAgentExecutor(auth_token="test-token") + + mock_message = mocker.MagicMock() + mock_message.role = "user" + mock_message.parts = [Part(root=TextPart(text="Hello"))] + mock_message.metadata = {} + + context = mocker.MagicMock(spec=RequestContext) + context.task_id = "task-123" + context.context_id = "ctx-456" + context.message = mock_message + context.get_user_input.return_value = "Query" + + event_queue = mocker.AsyncMock(spec=EventQueue) + task_updater = mocker.MagicMock() + task_updater.update_status = mocker.AsyncMock() + task_updater.event_queue = event_queue + + mock_context_store = mocker.AsyncMock() + mock_context_store.get.return_value = None + mocker.patch( + "app.endpoints.a2a._get_context_store", return_value=mock_context_store + ) + + mock_client = mocker.AsyncMock() + mock_client.models.list = mocker.AsyncMock( + return_value=ListModelsResponse.model_construct(data=[mocker.MagicMock()]) + ) + mocker.patch( + "app.endpoints.a2a.AsyncOgxClientHolder" + ).return_value.get_client.return_value = mock_client + + mock_responses_params = mocker.Mock() + mock_responses_params.model = "test-model" + mock_responses_params.conversation = "conv_x" + mocker.patch( + "app.endpoints.a2a.prepare_responses_params", + new=mocker.AsyncMock(return_value=mock_responses_params), + ) + + compaction_result = mocker.Mock() + compaction_result.params = mock_responses_params + mocker.patch( + "app.endpoints.a2a.apply_compaction_blocking", + new=mocker.AsyncMock(return_value=compaction_result), + ) + + mock_run_result = mocker.MagicMock() + mock_run_result.response.text = "Answer" + mock_usage = mocker.MagicMock() + mock_usage.input_tokens = 10 + mock_usage.output_tokens = 5 + mock_run_result.usage = mock_usage + + result_event = mocker.MagicMock(spec=AgentRunResultEvent) + result_event.result = mock_run_result + + async def _event_stream() -> Any: + yield result_event + + mock_stream_ctx = mocker.AsyncMock() + mock_stream_ctx.__aenter__ = mocker.AsyncMock(return_value=_event_stream()) + mock_stream_ctx.__aexit__ = mocker.AsyncMock(return_value=False) + mock_agent = mocker.MagicMock() + mock_agent.run_stream_events.return_value = mock_stream_ctx + mocker.patch("app.endpoints.a2a.build_agent", return_value=mock_agent) + + await executor._process_task_streaming( + context, task_updater, context.task_id, context.context_id + ) + + spans = exporter.get_finished_spans() + execute_spans = [s for s in spans if s.name == "a2a.execute"] + assert len(execute_spans) == 1 + span = execute_spans[0] + + event_names = [e.name for e in span.events] + assert "llm.inference.completed" in event_names + + @pytest.mark.asyncio + async def test_dispatch_span_attributes( # pylint: disable=too-many-locals + self, + mocker: MockerFixture, + setup_configuration: AppConfig, # pylint: disable=unused-argument + otel: tuple[Any, InMemorySpanExporter], + ) -> None: + """Test that a2a.dispatch span records rpc method, request id, and user_id.""" + tracer, exporter = otel + mocker.patch("app.endpoints.a2a.tracer", tracer) + + # Mock the A2A app to return a simple response + mock_a2a_app = mocker.AsyncMock() + + async def _mock_asgi(_scope: Any, _receive: Any, send: Any) -> None: + await send({"type": "http.response.start", "status": 200, "headers": []}) + await send( + { + "type": "http.response.body", + "body": b'{"jsonrpc":"2.0","id":"req-1","result":{}}', + } + ) + + mock_a2a_app.side_effect = _mock_asgi + mocker.patch( + "app.endpoints.a2a._create_a2a_app", + new=mocker.AsyncMock(return_value=mock_a2a_app), + ) + + # Build a mock request + rpc_body = b'{"jsonrpc":"2.0","method":"message/send","id":"req-1","params":{}}' + mock_request = mocker.MagicMock(spec=Request) + mock_request.method = "POST" + mock_request.url = mocker.MagicMock() + mock_request.url.path = "/a2a" + mock_request.body = mocker.AsyncMock(return_value=rpc_body) + mock_request.scope = { + "type": "http", + "method": "POST", + "path": "/a2a", + "headers": [], + } + mock_request.receive = mocker.AsyncMock() + mock_request.headers = {} + + await _handle_a2a_jsonrpc(mock_request, MOCK_AUTH, {}) + + spans = exporter.get_finished_spans() + dispatch_spans = [s for s in spans if s.name == "a2a.dispatch"] + assert len(dispatch_spans) == 1 + span = dispatch_spans[0] + attrs = dict(span.attributes or {}) + + assert attrs["a2a.rpc.method"] == "message/send" + assert attrs["a2a.request.id"].startswith("[hash:") + assert "user.id" in attrs + + event_names = [e.name for e in span.events] + assert "a2a.dispatch.start" in event_names + assert "a2a.dispatch.end" in event_names + + @pytest.mark.asyncio + async def test_dispatch_span_streaming_request( # pylint: disable=too-many-locals + self, + mocker: MockerFixture, + setup_configuration: AppConfig, # pylint: disable=unused-argument + otel: tuple[Any, InMemorySpanExporter], + ) -> None: + """Test that a2a.dispatch span is emitted for streaming requests.""" + tracer, exporter = otel + mocker.patch("app.endpoints.a2a.tracer", tracer) + + mock_a2a_app = mocker.AsyncMock() + + async def _mock_asgi(_scope: Any, _receive: Any, send: Any) -> None: + await send({"type": "http.response.start", "status": 200, "headers": []}) + await send( + { + "type": "http.response.body", + "body": b"data: chunk\n\n", + "more_body": False, + } + ) + + mock_a2a_app.side_effect = _mock_asgi + mocker.patch( + "app.endpoints.a2a._create_a2a_app", + new=mocker.AsyncMock(return_value=mock_a2a_app), + ) + + rpc_body = ( + b'{"jsonrpc":"2.0","method":"message/stream","id":"req-2","params":{}}' + ) + mock_request = mocker.MagicMock(spec=Request) + mock_request.method = "POST" + mock_request.url = mocker.MagicMock() + mock_request.url.path = "/a2a" + mock_request.body = mocker.AsyncMock(return_value=rpc_body) + mock_request.scope = { + "type": "http", + "method": "POST", + "path": "/a2a", + "headers": [], + } + mock_request.receive = mocker.AsyncMock() + mock_request.headers = {} + + response = await _handle_a2a_jsonrpc(mock_request, MOCK_AUTH, {}) + + # For streaming, consume the response generator to trigger app execution + assert hasattr(response, "body_iterator") + chunks = [] + async for chunk in response.body_iterator: + chunks.append(chunk) + + spans = exporter.get_finished_spans() + dispatch_spans = [s for s in spans if s.name == "a2a.dispatch"] + assert len(dispatch_spans) == 1 + span = dispatch_spans[0] + attrs = dict(span.attributes or {}) + + assert attrs["a2a.rpc.method"] == "message/stream" + assert attrs["a2a.request.id"].startswith("[hash:") + + event_names = [e.name for e in span.events] + assert "a2a.dispatch.start" in event_names + assert "a2a.dispatch.end" in event_names + + @pytest.mark.asyncio + async def test_execute_span_no_tool_calls( # pylint: disable=too-many-locals,too-many-statements + self, + mocker: MockerFixture, + setup_configuration: AppConfig, # pylint: disable=unused-argument + otel: tuple[Any, InMemorySpanExporter], + ) -> None: + """Test that tool attributes are absent when no tools are called.""" + tracer, exporter = otel + mocker.patch("app.endpoints.a2a.tracer", tracer) + + executor = A2AAgentExecutor(auth_token="test-token") + + mock_message = mocker.MagicMock() + mock_message.role = "user" + mock_message.parts = [Part(root=TextPart(text="Hello"))] + mock_message.metadata = {} + + context = mocker.MagicMock(spec=RequestContext) + context.task_id = "task-123" + context.context_id = "ctx-456" + context.message = mock_message + context.get_user_input.return_value = "Simple question" + + event_queue = mocker.AsyncMock(spec=EventQueue) + task_updater = mocker.MagicMock() + task_updater.update_status = mocker.AsyncMock() + task_updater.event_queue = event_queue + + mock_context_store = mocker.AsyncMock() + mock_context_store.get.return_value = None + mocker.patch( + "app.endpoints.a2a._get_context_store", return_value=mock_context_store + ) + + mock_client = mocker.AsyncMock() + mock_client.models.list = mocker.AsyncMock( + return_value=ListModelsResponse.model_construct(data=[mocker.MagicMock()]) + ) + mocker.patch( + "app.endpoints.a2a.AsyncOgxClientHolder" + ).return_value.get_client.return_value = mock_client + + mock_responses_params = mocker.Mock() + mock_responses_params.model = "test-model" + mock_responses_params.conversation = "conv_x" + mocker.patch( + "app.endpoints.a2a.prepare_responses_params", + new=mocker.AsyncMock(return_value=mock_responses_params), + ) + + compaction_result = mocker.Mock() + compaction_result.params = mock_responses_params + mocker.patch( + "app.endpoints.a2a.apply_compaction_blocking", + new=mocker.AsyncMock(return_value=compaction_result), + ) + + mock_run_result = mocker.MagicMock() + mock_run_result.response.text = "Simple answer" + mock_usage = mocker.MagicMock() + mock_usage.input_tokens = 5 + mock_usage.output_tokens = 3 + mock_run_result.usage = mock_usage + + result_event = mocker.MagicMock(spec=AgentRunResultEvent) + result_event.result = mock_run_result + + async def _event_stream() -> Any: + yield result_event + + mock_stream_ctx = mocker.AsyncMock() + mock_stream_ctx.__aenter__ = mocker.AsyncMock(return_value=_event_stream()) + mock_stream_ctx.__aexit__ = mocker.AsyncMock(return_value=False) + mock_agent = mocker.MagicMock() + mock_agent.run_stream_events.return_value = mock_stream_ctx + mocker.patch("app.endpoints.a2a.build_agent", return_value=mock_agent) + + await executor._process_task_streaming( + context, task_updater, context.task_id, context.context_id + ) + + spans = exporter.get_finished_spans() + execute_spans = [s for s in spans if s.name == "a2a.execute"] + assert len(execute_spans) == 1 + span = execute_spans[0] + attrs = dict(span.attributes or {}) + + assert "tool.calls.count" not in attrs + assert "tool.calls.names" not in attrs + + event_names = [e.name for e in span.events] + assert "tool.execution.completed" not in event_names