From c56957fc552ccced7e9d8984d8e78b759b4a3a33 Mon Sep 17 00:00:00 2001 From: Emmanuel Hugonnet Date: Wed, 9 Sep 2026 09:31:21 +0200 Subject: [PATCH] feat(opentelemetry): add streaming spans, metrics, and push tracing - Two-span model for server streaming: method span + closing span with events and link - Client streaming events recorded as span events on the parent span - gen_ai.client.operation.duration and gen_ai.agent.a2a.streaming.duration histograms - OpenTelemetryPushNotificationSenderDecorator with CLIENT spans for push delivery Signed-off-by: Emmanuel Hugonnet --- docs/content/dev/extra/opentelemetry.md | 105 +++++++- examples/helloworld/client/README.md | 4 +- .../helloworld/client/HelloWorldClient.java | 25 +- .../src/main/resources/application.properties | 1 + .../client/OpenTelemetryClientTransport.java | 145 +++++++--- .../OpenTelemetryClientTransportWrapper.java | 11 +- .../OpenTelemetryClientTransportTest.java | 103 +++++-- .../opentelemetry/A2AObservabilityNames.java | 9 + .../opentelemetry/integration-tests/pom.xml | 10 + .../opentelemetry/it/A2ATestRoutes.java | 111 ++++++++ .../it/TestAgentCardProducer.java | 2 +- .../src/main/resources/application.properties | 3 +- .../opentelemetry/it/A2ATestRoutesTest.java | 96 +++++++ .../sdk/extras/opentelemetry/it/BaseTest.java | 5 + .../it/OpenTelemetryA2ABaseTest.java | 153 +++++++++++ .../src/test/resources/application.properties | 3 +- ...emetryPushNotificationSenderDecorator.java | 146 ++++++++++ .../OpenTelemetryRequestHandlerDecorator.java | 201 ++++++++++++-- .../src/main/resources/META-INF/beans.xml | 3 +- ...ryPushNotificationSenderDecoratorTest.java | 187 +++++++++++++ ...nTelemetryRequestHandlerDecoratorTest.java | 251 +++++++++++++++++- 21 files changed, 1483 insertions(+), 91 deletions(-) create mode 100644 extras/opentelemetry/integration-tests/src/test/java/org/a2aproject/sdk/extras/opentelemetry/it/A2ATestRoutesTest.java create mode 100644 extras/opentelemetry/server/src/main/java/org/a2aproject/sdk/extras/opentelemetry/server/OpenTelemetryPushNotificationSenderDecorator.java create mode 100644 extras/opentelemetry/server/src/test/java/org/a2aproject/sdk/extras/opentelemetry/server/OpenTelemetryPushNotificationSenderDecoratorTest.java diff --git a/docs/content/dev/extra/opentelemetry.md b/docs/content/dev/extra/opentelemetry.md index a29e1b681..0093ffc9c 100644 --- a/docs/content/dev/extra/opentelemetry.md +++ b/docs/content/dev/extra/opentelemetry.md @@ -14,6 +14,7 @@ Adds distributed tracing and context propagation to A2A servers and clients usin - **Context Propagation**: OpenTelemetry trace context propagation across async operations - **Request/Response Logging**: Optional extraction of request and response data into spans - **Error Tracking**: Automatic error status and error type attributes on failures +- **Metrics**: `gen_ai.client.operation.duration` histogram on the client side; `gen_ai.agent.a2a.streaming.duration` histogram for full streaming lifecycle on the server side ## Modules @@ -65,6 +66,74 @@ Async agent execution (maintains trace context) Response (with trace headers) ``` +### Streaming Span Lifecycle + +Streaming methods (`message/stream` and `tasks/resubscribe`) follow a two-span model to correctly represent the full duration of a streaming response. + +**Span 1 — method span** is created when the request arrives and ended as soon as the publisher is returned to the transport layer. It captures request attributes and the initial outcome (success or synchronous error). + +**Span 2 — end span** (`-end`) is created when the stream terminates — either on completion or error. It carries a [link](https://opentelemetry.io/docs/concepts/signals/traces/#span-links) to Span 1 for correlation, and records every A2A event that was published as a [span event](https://opentelemetry.io/docs/concepts/signals/traces/#span-events). + +
+sequenceDiagram
+    participant C as Client
+    participant D as OTel Decorator
+    participant H as RequestHandler
+    participant W as WrappedPublisher
+
+    C->>D: onMessageSendStream(params)
+    D->>D: create Span₁ (SERVER)
+    D->>H: onMessageSendStream(params)
+    H-->>D: publisher
+    D->>D: capture Span₁.context
+    D->>D: Span₁.setStatus(OK) · Span₁.end()
+    D-->>C: WrappedPublisher
+
+    note over D: Span₁ closed immediately
+
+    loop for each streaming event
+        W-->>C: onNext(event)
+        W->>W: collect PendingEvent(name, attributes)
+    end
+
+    alt stream completes normally
+        W->>W: create Span₂ (SERVER, link→Span₁)
+        W->>W: addEvent per PendingEvent
+        W->>W: Span₂.setStatus(OK) · Span₂.end()
+        W-->>C: onComplete()
+    else stream errors
+        W->>W: create Span₂ (SERVER, link→Span₁)
+        W->>W: addEvent per PendingEvent
+        W->>W: Span₂.setStatus(ERROR) · Span₂.end()
+        W-->>C: onError(throwable)
+    end
+
+ +Each span event on Span₂ carries: + +| Attribute | Value | +|-----------|-------| +| `gen_ai.agent.a2a.streaming-event` | `toString()` of the published `StreamingEventKind` | +| `gen_ai.agent.a2a.status.code` | `OK` | + +### Metrics + +The server decorator records the following histogram instruments: + +| Instrument | Type | Unit | Description | +|------------|------|------|-------------| +| `gen_ai.agent.a2a.streaming.duration` | Histogram | `s` | Duration from stream initiation to last event, for `message/stream` and `tasks/resubscribe` | + +Attributes on `gen_ai.agent.a2a.streaming.duration`: + +| Attribute | Value | +|-----------|-------| +| `gen_ai.agent.a2a.operation.name` | A2A method name (`message/stream` or `tasks/resubscribe`) | +| `gen_ai.system` | `a2a` | +| `error.type` | `error` (only on stream failure) | + +Meter injection uses CDI `Instance`, so the histogram is silently skipped if no `Meter` bean is available in the container. + ### Context-Aware Async Executor > **Note:** The `AsyncManagedExecutorProducer` is provided by the **Quarkus reference server** ([`reference/common`](https://github.com/a2aproject/a2a-java/blob/main/reference/common/src/main/java/org/a2aproject/sdk/server/common/quarkus/AsyncManagedExecutorProducer.java)), not the OpenTelemetry module. It is documented here because it enables context propagation (including trace context) across async boundaries. @@ -106,6 +175,19 @@ The following attributes are automatically added to spans (defined in `A2AObserv | `gen_ai.agent.a2a.response` | Full response data (only if extraction enabled) | | `error.type` | Error message (on failures) | +### Push Notification Delivery Tracing + +When the server delivers a push notification to a client webhook, `OpenTelemetryPushNotificationSenderDecorator` creates a `CLIENT` span named `SendPushNotification`. This span is independent of the original request span but carries the `gen_ai.agent.a2a.task_id` attribute, allowing correlation by task ID in your observability tool. + +| Attribute | Description | +|-----------|-------------| +| `gen_ai.agent.a2a.operation.name` | `SendPushNotification` | +| `gen_ai.agent.a2a.task_id` | Task identifier | +| `gen_ai.agent.a2a.context_id` | Context identifier (when available) | +| `gen_ai.agent.a2a.push_notification.event_kind` | The type of streaming event delivered (e.g. `TaskStatusUpdateEvent`) | +| `gen_ai.agent.a2a.response` | Full event payload (only if extraction enabled) | +| `error.type` | Error message (on delivery failure) | + ### Request/Response Extraction Enable request and response data extraction in spans using JVM system properties: @@ -124,7 +206,7 @@ Enable request and response data extraction in spans using JVM system properties ### Instrumentation -Adds OpenTelemetry spans to A2A client operations: +Adds OpenTelemetry spans and metrics to A2A client operations: ```xml @@ -133,6 +215,27 @@ Adds OpenTelemetry spans to A2A client operations: ``` +To also enable `gen_ai.client.operation.duration` metrics, pass a `Meter` via the transport config: + +```java +config.setParameters(Map.of( + OpenTelemetryClientTransportWrapper.OTEL_TRACER_KEY, openTelemetry.getTracer("my-service"), + OpenTelemetryClientTransportWrapper.OTEL_METER_KEY, openTelemetry.getMeter("my-service") +)); +``` + +| Instrument | Type | Unit | Description | +|------------|------|------|-------------| +| `gen_ai.client.operation.duration` | Histogram | `s` | Duration of each A2A client operation call | + +Attributes on `gen_ai.client.operation.duration`: + +| Attribute | Value | +|-----------|-------| +| `gen_ai.agent.a2a.operation.name` | A2A method name (e.g. `message/send`) | +| `gen_ai.system` | `a2a` | +| `error.type` | `error` (only on failure) | + ### Context Propagation For trace context propagation across async client boundaries: diff --git a/examples/helloworld/client/README.md b/examples/helloworld/client/README.md index b186f7ec7..d0673779f 100644 --- a/examples/helloworld/client/README.md +++ b/examples/helloworld/client/README.md @@ -161,11 +161,11 @@ The Java client (`HelloWorldClient.java`) performs the following actions: ## Enable OpenTelemetry (Optional) -The client includes support for distributed tracing with OpenTelemetry. To enable it: +The client includes support for distributed tracing and metrics with OpenTelemetry. To enable it: ### Prerequisites -**IMPORTANT**: The client expects an OpenTelemetry collector to be ready and accepting traces. You have two options: +**IMPORTANT**: The client expects an OpenTelemetry collector to be ready and accepting traces and metrics. You have two options: #### Option 1: Use the Java Server Example (Recommended) diff --git a/examples/helloworld/client/src/main/java/org/a2aproject/sdk/examples/helloworld/client/HelloWorldClient.java b/examples/helloworld/client/src/main/java/org/a2aproject/sdk/examples/helloworld/client/HelloWorldClient.java index 0c24d08c2..46a80ecb2 100644 --- a/examples/helloworld/client/src/main/java/org/a2aproject/sdk/examples/helloworld/client/HelloWorldClient.java +++ b/examples/helloworld/client/src/main/java/org/a2aproject/sdk/examples/helloworld/client/HelloWorldClient.java @@ -1,5 +1,6 @@ package org.a2aproject.sdk.examples.helloworld.client; +import static org.a2aproject.sdk.extras.opentelemetry.client.OpenTelemetryClientTransportWrapper.OTEL_METER_KEY; import static org.a2aproject.sdk.extras.opentelemetry.client.OpenTelemetryClientTransportWrapper.OTEL_TRACER_KEY; import static org.a2aproject.sdk.extras.opentelemetry.client.propagation.OpenTelemetryClientPropagatorTransportWrapper.OTEL_OPEN_TELEMETRY_KEY; @@ -34,8 +35,11 @@ import io.grpc.ManagedChannelBuilder; import io.opentelemetry.api.trace.propagation.W3CTraceContextPropagator; import io.opentelemetry.context.propagation.ContextPropagators; +import io.opentelemetry.exporter.otlp.metrics.OtlpGrpcMetricExporter; import io.opentelemetry.exporter.otlp.trace.OtlpGrpcSpanExporter; import io.opentelemetry.sdk.OpenTelemetrySdk; +import io.opentelemetry.sdk.metrics.SdkMeterProvider; +import io.opentelemetry.sdk.metrics.export.PeriodicMetricReader; import io.opentelemetry.sdk.resources.Resource; import io.opentelemetry.sdk.trace.SdkTracerProvider; import io.opentelemetry.sdk.trace.export.BatchSpanProcessor; @@ -136,20 +140,32 @@ public static void main(String[] args) { } static OpenTelemetrySdk initOpenTelemetry() { + Resource resource = Resource.getDefault().toBuilder() + .put("service.version", "1.0") + .put("service.name", "helloworld-client") + .build(); + SdkTracerProvider sdkTracerProvider = SdkTracerProvider.builder() .addSpanProcessor(BatchSpanProcessor.builder( OtlpGrpcSpanExporter.builder() .setEndpoint("http://localhost:5317") .build() ).build()) - .setResource(Resource.getDefault().toBuilder() - .put("service.version", "1.0") - .put("service.name", "helloworld-client") - .build()) + .setResource(resource) + .build(); + + SdkMeterProvider sdkMeterProvider = SdkMeterProvider.builder() + .registerMetricReader(PeriodicMetricReader.builder( + OtlpGrpcMetricExporter.builder() + .setEndpoint("http://localhost:5317") + .build() + ).build()) + .setResource(resource) .build(); return OpenTelemetrySdk.builder() .setTracerProvider(sdkTracerProvider) + .setMeterProvider(sdkMeterProvider) .setPropagators(ContextPropagators.create(W3CTraceContextPropagator.getInstance())) .build(); } @@ -186,6 +202,7 @@ private static void updateTransportConfig(ClientTransportConfig transportConfig, if (openTelemetrySdk != null) { Map parameters = new HashMap<>(transportConfig.getParameters()); parameters.put(OTEL_TRACER_KEY, openTelemetrySdk.getTracer("helloworld-client")); + parameters.put(OTEL_METER_KEY, openTelemetrySdk.getMeter("helloworld-client")); parameters.put(OTEL_OPEN_TELEMETRY_KEY, openTelemetrySdk); transportConfig.setParameters(parameters); } diff --git a/examples/helloworld/server/src/main/resources/application.properties b/examples/helloworld/server/src/main/resources/application.properties index 4edbf121b..8e406688a 100644 --- a/examples/helloworld/server/src/main/resources/application.properties +++ b/examples/helloworld/server/src/main/resources/application.properties @@ -5,6 +5,7 @@ quarkus.agentcard.protocol=JSONRPC # OpenTelemetry configuration quarkus.otel.sdk.disabled=false +quarkus.otel.metrics.enabled=true quarkus.observability.lgtm.grafana-port=3001 quarkus.observability.lgtm.otel-grpc-port=5317 quarkus.observability.lgtm.otel-http-port=5318 diff --git a/extras/opentelemetry/client/src/main/java/org/a2aproject/sdk/extras/opentelemetry/client/OpenTelemetryClientTransport.java b/extras/opentelemetry/client/src/main/java/org/a2aproject/sdk/extras/opentelemetry/client/OpenTelemetryClientTransport.java index 574253978..5439de7ad 100644 --- a/extras/opentelemetry/client/src/main/java/org/a2aproject/sdk/extras/opentelemetry/client/OpenTelemetryClientTransport.java +++ b/extras/opentelemetry/client/src/main/java/org/a2aproject/sdk/extras/opentelemetry/client/OpenTelemetryClientTransport.java @@ -1,7 +1,9 @@ package org.a2aproject.sdk.extras.opentelemetry.client; +import static org.a2aproject.sdk.extras.opentelemetry.A2AObservabilityNames.ERROR_TYPE; import static org.a2aproject.sdk.extras.opentelemetry.A2AObservabilityNames.EXTRACT_REQUEST_SYS_PROPERTY; import static org.a2aproject.sdk.extras.opentelemetry.A2AObservabilityNames.EXTRACT_RESPONSE_SYS_PROPERTY; +import static org.a2aproject.sdk.extras.opentelemetry.A2AObservabilityNames.GENAI_CLIENT_OPERATION_DURATION; import static org.a2aproject.sdk.extras.opentelemetry.A2AObservabilityNames.GENAI_CONFIG_ID; import static org.a2aproject.sdk.extras.opentelemetry.A2AObservabilityNames.GENAI_CONTEXT_ID; import static org.a2aproject.sdk.extras.opentelemetry.A2AObservabilityNames.GENAI_EXTENSIONS; @@ -11,8 +13,24 @@ import static org.a2aproject.sdk.extras.opentelemetry.A2AObservabilityNames.GENAI_REQUEST; import static org.a2aproject.sdk.extras.opentelemetry.A2AObservabilityNames.GENAI_RESPONSE; import static org.a2aproject.sdk.extras.opentelemetry.A2AObservabilityNames.GENAI_ROLE; +import static org.a2aproject.sdk.extras.opentelemetry.A2AObservabilityNames.GENAI_SYSTEM; +import static org.a2aproject.sdk.extras.opentelemetry.A2AObservabilityNames.GENAI_SYSTEM_VALUE; import static org.a2aproject.sdk.extras.opentelemetry.A2AObservabilityNames.GENAI_TASK_ID; +import io.opentelemetry.api.common.Attributes; +import io.opentelemetry.api.common.AttributesBuilder; +import io.opentelemetry.api.metrics.DoubleHistogram; +import io.opentelemetry.api.metrics.Meter; +import io.opentelemetry.api.trace.Span; +import io.opentelemetry.api.trace.SpanBuilder; +import io.opentelemetry.api.trace.SpanKind; +import io.opentelemetry.api.trace.StatusCode; +import io.opentelemetry.api.trace.Tracer; +import io.opentelemetry.context.Scope; +import java.util.HashMap; +import java.util.Map; +import java.util.function.Consumer; +import java.util.stream.Collectors; import org.a2aproject.sdk.client.transport.spi.ClientTransport; import org.a2aproject.sdk.client.transport.spi.interceptors.ClientCallContext; import org.a2aproject.sdk.jsonrpc.common.wrappers.ListTasksResult; @@ -33,27 +51,29 @@ import org.a2aproject.sdk.spec.TaskIdParams; import org.a2aproject.sdk.spec.TaskPushNotificationConfig; import org.a2aproject.sdk.spec.TaskQueryParams; -import io.opentelemetry.api.trace.Span; -import io.opentelemetry.api.trace.SpanBuilder; -import io.opentelemetry.api.trace.SpanContext; -import io.opentelemetry.api.trace.SpanKind; -import io.opentelemetry.api.trace.StatusCode; -import io.opentelemetry.api.trace.Tracer; -import io.opentelemetry.context.Scope; -import java.util.HashMap; -import java.util.Map; -import java.util.function.Consumer; -import java.util.stream.Collectors; import org.jspecify.annotations.Nullable; public class OpenTelemetryClientTransport implements ClientTransport { private final Tracer tracer; private final ClientTransport delegate; + @Nullable private final DoubleHistogram operationDurationHistogram; public OpenTelemetryClientTransport(ClientTransport delegate, Tracer tracer) { + this(delegate, tracer, null); + } + + /** + * @param meter optional meter for recording {@code gen_ai.client.operation.duration}; null disables metrics + */ + public OpenTelemetryClientTransport(ClientTransport delegate, Tracer tracer, @Nullable Meter meter) { this.delegate = delegate; this.tracer = tracer; + this.operationDurationHistogram = meter == null ? null + : meter.histogramBuilder(GENAI_CLIENT_OPERATION_DURATION) + .setUnit("s") + .setDescription("Duration of A2A client operations") + .build(); } private boolean extractRequest() { @@ -67,6 +87,7 @@ private boolean extractResponse() { @Override public EventKind sendMessage(MessageSendParams request, @Nullable ClientCallContext context) throws A2AClientException { ClientCallContext clientContext = createContext(context); + long startNanos = System.nanoTime(); SpanBuilder spanBuilder = tracer.spanBuilder(A2AMethods.SEND_MESSAGE_METHOD).setSpanKind(SpanKind.CLIENT); spanBuilder.setAttribute(GENAI_OPERATION_NAME, A2AMethods.SEND_MESSAGE_METHOD); if (request.message() != null) { @@ -91,6 +112,7 @@ public EventKind sendMessage(MessageSendParams request, @Nullable ClientCallCont spanBuilder.setAttribute(GENAI_REQUEST, request.toString()); } Span span = spanBuilder.startSpan(); + boolean success = false; try (Scope scope = span.makeCurrent()) { EventKind result = delegate.sendMessage(request, clientContext); if (result != null && extractResponse()) { @@ -99,12 +121,14 @@ public EventKind sendMessage(MessageSendParams request, @Nullable ClientCallCont if (result != null) { span.setStatus(StatusCode.OK); } + success = true; return result; } catch (Exception ex) { span.setStatus(StatusCode.ERROR, ex.getMessage()); throw ex; } finally { span.end(); + recordOperationDuration(startNanos, A2AMethods.SEND_MESSAGE_METHOD, success); } } @@ -112,6 +136,7 @@ public EventKind sendMessage(MessageSendParams request, @Nullable ClientCallCont public void sendMessageStreaming(MessageSendParams request, Consumer eventConsumer, Consumer errorConsumer, @Nullable ClientCallContext context) throws A2AClientException { ClientCallContext clientContext = createContext(context); + long startNanos = System.nanoTime(); SpanBuilder spanBuilder = tracer.spanBuilder(A2AMethods.SEND_STREAMING_MESSAGE_METHOD).setSpanKind(SpanKind.CLIENT); spanBuilder.setAttribute(GENAI_OPERATION_NAME, A2AMethods.SEND_STREAMING_MESSAGE_METHOD); if (request.message() != null) { @@ -136,25 +161,29 @@ public void sendMessageStreaming(MessageSendParams request, Consumer eventConsumer, Consumer errorConsumer, @Nullable ClientCallContext context) throws A2AClientException { ClientCallContext clientContext = createContext(context); + long startNanos = System.nanoTime(); SpanBuilder spanBuilder = tracer.spanBuilder(A2AMethods.SUBSCRIBE_TO_TASK_METHOD).setSpanKind(SpanKind.CLIENT); spanBuilder.setAttribute(GENAI_OPERATION_NAME, A2AMethods.SUBSCRIBE_TO_TASK_METHOD); if (request.id() != null) { @@ -378,28 +435,33 @@ public void subscribeToTask(TaskIdParams request, Consumer e spanBuilder.setAttribute(GENAI_REQUEST, request.toString()); } Span span = spanBuilder.startSpan(); + boolean success = false; try (Scope scope = span.makeCurrent()) { delegate.subscribeToTask( request, - new OpenTelemetryEventConsumer(A2AMethods.SUBSCRIBE_TO_TASK_METHOD + "-event", eventConsumer, tracer, span.getSpanContext()), - new OpenTelemetryErrorConsumer(A2AMethods.SUBSCRIBE_TO_TASK_METHOD + "-error", errorConsumer, tracer, span.getSpanContext()), + new OpenTelemetryEventConsumer(A2AMethods.SUBSCRIBE_TO_TASK_METHOD + "-event", eventConsumer, span), + new OpenTelemetryErrorConsumer(A2AMethods.SUBSCRIBE_TO_TASK_METHOD + "-error", errorConsumer, span), clientContext ); span.setStatus(StatusCode.OK); + success = true; } catch (Exception ex) { span.setStatus(StatusCode.ERROR, ex.getMessage()); throw ex; } finally { span.end(); + recordOperationDuration(startNanos, A2AMethods.SUBSCRIBE_TO_TASK_METHOD, success); } } @Override public AgentCard getExtendedAgentCard(GetExtendedAgentCardParams params, @Nullable ClientCallContext context) throws A2AClientException { ClientCallContext clientContext = createContext(context); + long startNanos = System.nanoTime(); SpanBuilder spanBuilder = tracer.spanBuilder(A2AMethods.GET_EXTENDED_AGENT_CARD_METHOD).setSpanKind(SpanKind.CLIENT); spanBuilder.setAttribute(GENAI_OPERATION_NAME, A2AMethods.GET_EXTENDED_AGENT_CARD_METHOD); Span span = spanBuilder.startSpan(); + boolean success = false; try (Scope scope = span.makeCurrent()) { AgentCard result = delegate.getExtendedAgentCard(params, clientContext); if (result != null && extractResponse()) { @@ -408,15 +470,31 @@ public AgentCard getExtendedAgentCard(GetExtendedAgentCardParams params, @Nullab if (result != null) { span.setStatus(StatusCode.OK); } + success = true; return result; } catch (Exception ex) { span.setStatus(StatusCode.ERROR, ex.getMessage()); throw ex; } finally { span.end(); + recordOperationDuration(startNanos, A2AMethods.GET_EXTENDED_AGENT_CARD_METHOD, success); } } + private void recordOperationDuration(long startNanos, String operationName, boolean success) { + if (operationDurationHistogram == null) { + return; + } + double seconds = (System.nanoTime() - startNanos) / 1_000_000_000.0; + AttributesBuilder builder = Attributes.builder() + .put(GENAI_OPERATION_NAME, operationName) + .put(GENAI_SYSTEM, GENAI_SYSTEM_VALUE); + if (!success) { + builder.put(ERROR_TYPE, "error"); + } + operationDurationHistogram.record(seconds, builder.build()); + } + private ClientCallContext createContext(@Nullable ClientCallContext context) { if (context == null) { return new ClientCallContext(Map.of(), new HashMap<>()); @@ -432,29 +510,24 @@ public void close() { private static class OpenTelemetryEventConsumer implements Consumer { private final Consumer delegate; - private final Tracer tracer; - private final SpanContext context; + private final Span span; private final String name; - public OpenTelemetryEventConsumer(String name, Consumer delegate, Tracer tracer, SpanContext context) { + public OpenTelemetryEventConsumer(String name, Consumer delegate, Span span) { this.delegate = delegate; - this.tracer = tracer; - this.context = context; + this.span = span; this.name = name; } @Override public void accept(StreamingEventKind t) { - SpanBuilder spanBuilder = tracer.spanBuilder(name) - .setSpanKind(SpanKind.CLIENT); - spanBuilder.setAttribute("gen_ai.agent.a2a.streaming-event", t.toString()); - spanBuilder.addLink(context); - Span span = spanBuilder.startSpan(); + AttributesBuilder builder = Attributes.builder(); + builder.put("gen_ai.agent.a2a.streaming-event", t.toString()); try { delegate.accept(t); - span.setStatus(StatusCode.OK); + builder.put("gen_ai.agent.a2a.status.code", StatusCode.OK.name()); } finally { - span.end(); + span.addEvent(name, builder.build()); } } } @@ -462,14 +535,12 @@ public void accept(StreamingEventKind t) { private static class OpenTelemetryErrorConsumer implements Consumer { private final Consumer delegate; - private final Tracer tracer; - private final SpanContext context; + private final Span span; private final String name; - public OpenTelemetryErrorConsumer(String name, Consumer delegate, Tracer tracer, SpanContext context) { + public OpenTelemetryErrorConsumer(String name, Consumer delegate, Span span) { this.delegate = delegate; - this.tracer = tracer; - this.context = context; + this.span = span; this.name = name; } @@ -478,15 +549,13 @@ public void accept(Throwable t) { if (t == null) { return; } - SpanBuilder spanBuilder = tracer.spanBuilder(name) - .setSpanKind(SpanKind.CLIENT); - spanBuilder.addLink(context); - Span span = spanBuilder.startSpan(); + AttributesBuilder builder = Attributes.builder(); + builder.put("gen_ai.agent.a2a.status.code", StatusCode.ERROR.name()); + builder.put("gen_ai.agent.a2a.status.description", t.getMessage()); try { - span.setStatus(StatusCode.ERROR, t.getMessage()); delegate.accept(t); } finally { - span.end(); + span.addEvent(name, builder.build()); } } } diff --git a/extras/opentelemetry/client/src/main/java/org/a2aproject/sdk/extras/opentelemetry/client/OpenTelemetryClientTransportWrapper.java b/extras/opentelemetry/client/src/main/java/org/a2aproject/sdk/extras/opentelemetry/client/OpenTelemetryClientTransportWrapper.java index 5b329c8f1..0e4bb07a6 100644 --- a/extras/opentelemetry/client/src/main/java/org/a2aproject/sdk/extras/opentelemetry/client/OpenTelemetryClientTransportWrapper.java +++ b/extras/opentelemetry/client/src/main/java/org/a2aproject/sdk/extras/opentelemetry/client/OpenTelemetryClientTransportWrapper.java @@ -3,6 +3,7 @@ import org.a2aproject.sdk.client.transport.spi.ClientTransport; import org.a2aproject.sdk.client.transport.spi.ClientTransportConfig; import org.a2aproject.sdk.client.transport.spi.ClientTransportWrapper; +import io.opentelemetry.api.metrics.Meter; import io.opentelemetry.api.trace.Tracer; /** @@ -26,11 +27,19 @@ public class OpenTelemetryClientTransportWrapper implements ClientTransportWrapp */ public static final String OTEL_TRACER_KEY = "org.a2aproject.sdk.extras.opentelemetry.Tracer"; + /** + * Configuration key for the OpenTelemetry Meter instance. + * Value must be of type {@link Meter}. + */ + public static final String OTEL_METER_KEY = "org.a2aproject.sdk.extras.opentelemetry.Meter"; + @Override public ClientTransport wrap(ClientTransport transport, ClientTransportConfig config) { Object tracerObj = config.getParameters().get(OTEL_TRACER_KEY); if (tracerObj != null && tracerObj instanceof Tracer tracer) { - return new OpenTelemetryClientTransport(transport, tracer); + Object meterObj = config.getParameters().get(OTEL_METER_KEY); + Meter meter = (meterObj instanceof Meter m) ? m : null; + return new OpenTelemetryClientTransport(transport, tracer, meter); } // No tracer configured, return unwrapped transport return transport; diff --git a/extras/opentelemetry/client/src/test/java/org/a2aproject/sdk/extras/opentelemetry/client/OpenTelemetryClientTransportTest.java b/extras/opentelemetry/client/src/test/java/org/a2aproject/sdk/extras/opentelemetry/client/OpenTelemetryClientTransportTest.java index fd1e7d992..1b80229a5 100644 --- a/extras/opentelemetry/client/src/test/java/org/a2aproject/sdk/extras/opentelemetry/client/OpenTelemetryClientTransportTest.java +++ b/extras/opentelemetry/client/src/test/java/org/a2aproject/sdk/extras/opentelemetry/client/OpenTelemetryClientTransportTest.java @@ -25,6 +25,11 @@ import org.a2aproject.sdk.spec.TaskIdParams; import org.a2aproject.sdk.spec.TaskPushNotificationConfig; import org.a2aproject.sdk.spec.TaskQueryParams; +import io.opentelemetry.api.common.AttributeKey; +import io.opentelemetry.api.common.Attributes; +import io.opentelemetry.api.metrics.DoubleHistogram; +import io.opentelemetry.api.metrics.DoubleHistogramBuilder; +import io.opentelemetry.api.metrics.Meter; import io.opentelemetry.api.trace.Span; import io.opentelemetry.api.trace.SpanBuilder; import io.opentelemetry.api.trace.SpanContext; @@ -33,6 +38,7 @@ import io.opentelemetry.api.trace.Tracer; import io.opentelemetry.context.Scope; import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Nested; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; import org.mockito.ArgumentCaptor; @@ -364,14 +370,6 @@ void testEventConsumer_ThroughSendMessageStreaming() throws A2AClientException { MessageSendParams request = mock(MessageSendParams.class); when(request.toString()).thenReturn("request-string"); - SpanBuilder eventSpanBuilder = mock(SpanBuilder.class); - Span eventSpan = mock(Span.class); - when(tracer.spanBuilder(A2AMethods.SEND_STREAMING_MESSAGE_METHOD + "-event")).thenReturn(eventSpanBuilder); - when(eventSpanBuilder.setSpanKind(any(SpanKind.class))).thenReturn(eventSpanBuilder); - when(eventSpanBuilder.setAttribute(anyString(), anyString())).thenReturn(eventSpanBuilder); - when(eventSpanBuilder.addLink(any(SpanContext.class))).thenReturn(eventSpanBuilder); - when(eventSpanBuilder.startSpan()).thenReturn(eventSpan); - ArgumentCaptor> eventConsumerCaptor = ArgumentCaptor.forClass(Consumer.class); Consumer originalConsumer = mock(Consumer.class); @@ -388,9 +386,7 @@ void testEventConsumer_ThroughSendMessageStreaming() throws A2AClientException { eventConsumerCaptor.getValue().accept(event); - verify(tracer).spanBuilder(A2AMethods.SEND_STREAMING_MESSAGE_METHOD + "-event"); - verify(eventSpan).setStatus(StatusCode.OK); - verify(eventSpan).end(); + verify(span).addEvent(eq(A2AMethods.SEND_STREAMING_MESSAGE_METHOD + "-event"), any(io.opentelemetry.api.common.Attributes.class)); verify(originalConsumer).accept(event); } @@ -399,13 +395,6 @@ void testErrorConsumer_ThroughSendMessageStreaming() throws A2AClientException { MessageSendParams request = mock(MessageSendParams.class); when(request.toString()).thenReturn("request-string"); - SpanBuilder errorSpanBuilder = mock(SpanBuilder.class); - Span errorSpan = mock(Span.class); - when(tracer.spanBuilder(A2AMethods.SEND_STREAMING_MESSAGE_METHOD + "-error")).thenReturn(errorSpanBuilder); - when(errorSpanBuilder.setSpanKind(any(SpanKind.class))).thenReturn(errorSpanBuilder); - when(errorSpanBuilder.addLink(any(SpanContext.class))).thenReturn(errorSpanBuilder); - when(errorSpanBuilder.startSpan()).thenReturn(errorSpan); - ArgumentCaptor> errorConsumerCaptor = ArgumentCaptor.forClass(Consumer.class); Consumer originalConsumer = mock(Consumer.class); @@ -417,9 +406,7 @@ void testErrorConsumer_ThroughSendMessageStreaming() throws A2AClientException { errorConsumerCaptor.getValue().accept(error); - verify(tracer).spanBuilder(A2AMethods.SEND_STREAMING_MESSAGE_METHOD + "-error"); - verify(errorSpan).setStatus(StatusCode.ERROR, "Test error"); - verify(errorSpan).end(); + verify(span).addEvent(eq(A2AMethods.SEND_STREAMING_MESSAGE_METHOD + "-error"), any(io.opentelemetry.api.common.Attributes.class)); verify(originalConsumer).accept(error); } @@ -473,6 +460,80 @@ void testResubscribe_ThrowsException() throws A2AClientException { verify(span).end(); } + @Nested + class MetricsTests { + + @Mock + private Meter meter; + + @Mock + private DoubleHistogramBuilder histogramBuilder; + + @Mock + private DoubleHistogram histogram; + + @BeforeEach + void setUpMeter() { + lenient().when(meter.histogramBuilder(anyString())).thenReturn(histogramBuilder); + lenient().when(histogramBuilder.setUnit(anyString())).thenReturn(histogramBuilder); + lenient().when(histogramBuilder.setDescription(anyString())).thenReturn(histogramBuilder); + lenient().when(histogramBuilder.build()).thenReturn(histogram); + } + + @Test + void sendMessage_recordsOperationDuration() throws A2AClientException { + OpenTelemetryClientTransport transportWithMeter = new OpenTelemetryClientTransport(delegate, tracer, meter); + MessageSendParams request = mock(MessageSendParams.class); + EventKind expectedResult = mock(EventKind.class); + when(request.toString()).thenReturn("request-string"); + when(expectedResult.toString()).thenReturn("response-string"); + when(delegate.sendMessage(eq(request), any(ClientCallContext.class))).thenReturn(expectedResult); + + transportWithMeter.sendMessage(request, context); + + verify(histogram).record(anyDouble(), any(io.opentelemetry.api.common.Attributes.class)); + } + + @Test + void sendMessage_noHistogramWhenNoMeter() throws A2AClientException { + // transport without meter — should not throw, histogram never invoked + MessageSendParams request = mock(MessageSendParams.class); + EventKind expectedResult = mock(EventKind.class); + when(request.toString()).thenReturn("request-string"); + when(expectedResult.toString()).thenReturn("response-string"); + when(delegate.sendMessage(eq(request), any(ClientCallContext.class))).thenReturn(expectedResult); + + transport.sendMessage(request, context); + + verify(histogram, never()).record(anyDouble(), any(io.opentelemetry.api.common.Attributes.class)); + } + } + + @Test + void testErrorConsumer_emitsStatusAttributesWithoutStreamingEventAttribute() throws A2AClientException { + MessageSendParams request = mock(MessageSendParams.class); + when(request.toString()).thenReturn("request-string"); + Consumer originalConsumer = mock(Consumer.class); + + transport.sendMessageStreaming(request, mock(Consumer.class), originalConsumer, context); + + ArgumentCaptor> errorConsumerCaptor = ArgumentCaptor.forClass(Consumer.class); + verify(delegate).sendMessageStreaming(eq(request), any(Consumer.class), + errorConsumerCaptor.capture(), any(ClientCallContext.class)); + + RuntimeException error = new RuntimeException("stream broke"); + errorConsumerCaptor.getValue().accept(error); + + ArgumentCaptor attrsCaptor = ArgumentCaptor.forClass(Attributes.class); + verify(span).addEvent(eq(A2AMethods.SEND_STREAMING_MESSAGE_METHOD + "-error"), attrsCaptor.capture()); + + Attributes attrs = attrsCaptor.getValue(); + assertEquals(StatusCode.ERROR.name(), attrs.get(AttributeKey.stringKey("gen_ai.agent.a2a.status.code"))); + assertEquals("stream broke", attrs.get(AttributeKey.stringKey("gen_ai.agent.a2a.status.description"))); + assertNull(attrs.get(AttributeKey.stringKey("gen_ai.agent.a2a.streaming-event")), + "streaming-event must not be set on an error span event"); + } + @Test void testSendMessageStreaming_ThrowsException() throws A2AClientException { MessageSendParams request = mock(MessageSendParams.class); diff --git a/extras/opentelemetry/common/src/main/java/org/a2aproject/sdk/extras/opentelemetry/A2AObservabilityNames.java b/extras/opentelemetry/common/src/main/java/org/a2aproject/sdk/extras/opentelemetry/A2AObservabilityNames.java index 4bd74e99b..964b7b2d0 100644 --- a/extras/opentelemetry/common/src/main/java/org/a2aproject/sdk/extras/opentelemetry/A2AObservabilityNames.java +++ b/extras/opentelemetry/common/src/main/java/org/a2aproject/sdk/extras/opentelemetry/A2AObservabilityNames.java @@ -20,4 +20,13 @@ public interface A2AObservabilityNames { String GENAI_RESPONSE = GENAI_PREFIX + ".response"; // gen_ai.output.messages ? String GENAI_ROLE = GENAI_PREFIX + ".role"; String GENAI_TASK_ID = GENAI_PREFIX + ".task_id"; + + // Push notification + String PUSH_NOTIFICATION_EVENT_KIND = GENAI_PREFIX + ".push_notification.event_kind"; + + // Metrics + String GENAI_CLIENT_OPERATION_DURATION = "gen_ai.client.operation.duration"; + String GENAI_STREAMING_DURATION = GENAI_PREFIX + ".streaming.duration"; + String GENAI_SYSTEM = "gen_ai.system"; + String GENAI_SYSTEM_VALUE = "a2a"; } diff --git a/extras/opentelemetry/integration-tests/pom.xml b/extras/opentelemetry/integration-tests/pom.xml index 3daff9239..686a75fa0 100644 --- a/extras/opentelemetry/integration-tests/pom.xml +++ b/extras/opentelemetry/integration-tests/pom.xml @@ -83,6 +83,16 @@ awaitility test + + org.mockito + mockito-core + test + + + org.mockito + mockito-junit-jupiter + test + diff --git a/extras/opentelemetry/integration-tests/src/main/java/org/a2aproject/sdk/extras/opentelemetry/it/A2ATestRoutes.java b/extras/opentelemetry/integration-tests/src/main/java/org/a2aproject/sdk/extras/opentelemetry/it/A2ATestRoutes.java index 4c37b5c5d..69c0ba2d5 100644 --- a/extras/opentelemetry/integration-tests/src/main/java/org/a2aproject/sdk/extras/opentelemetry/it/A2ATestRoutes.java +++ b/extras/opentelemetry/integration-tests/src/main/java/org/a2aproject/sdk/extras/opentelemetry/it/A2ATestRoutes.java @@ -13,9 +13,15 @@ import org.a2aproject.sdk.spec.Task; import org.a2aproject.sdk.spec.TaskArtifactUpdateEvent; import org.a2aproject.sdk.spec.TaskStatusUpdateEvent; +import io.opentelemetry.api.OpenTelemetry; import io.opentelemetry.api.trace.Span; import io.opentelemetry.api.trace.Tracer; import io.opentelemetry.context.Scope; +import io.opentelemetry.sdk.OpenTelemetrySdk; +import io.opentelemetry.sdk.metrics.data.HistogramData; +import io.opentelemetry.sdk.metrics.data.MetricData; +import io.opentelemetry.sdk.metrics.data.MetricDataType; +import io.opentelemetry.sdk.testing.exporter.InMemoryMetricExporter; import io.opentelemetry.sdk.testing.exporter.InMemorySpanExporter; import io.opentelemetry.sdk.trace.data.SpanData; import io.vertx.ext.web.Router; @@ -27,6 +33,7 @@ import jakarta.inject.Inject; import jakarta.inject.Singleton; import java.util.List; +import java.util.concurrent.atomic.AtomicInteger; import java.util.stream.Collectors; /** @@ -40,10 +47,16 @@ public class A2ATestRoutes { private static final String TEXT_PLAIN = "text/plain"; private static final Gson gson = new GsonBuilder().create(); + private final AtomicInteger webhookCallCount = new AtomicInteger(0); + @Inject TestUtilsBean testUtilsBean; @Inject InMemorySpanExporter inMemorySpanExporter; + @Inject + InMemoryMetricExporter inMemoryMetricExporter; + @Inject + OpenTelemetry openTelemetry; @Inject Tracer tracer; @@ -84,9 +97,25 @@ void setupRoutes(@Observes Router router) { .produces(APPLICATION_JSON) .handler(ctx -> exportSpans(ctx)); + router.get("/export-metrics") + .produces(APPLICATION_JSON) + .handler(ctx -> exportMetrics(ctx)); + router.get("/reset") .produces(TEXT_PLAIN) .handler(ctx -> reset(ctx)); + + router.get("/reset-metrics") + .produces(TEXT_PLAIN) + .handler(ctx -> resetMetrics(ctx)); + + router.post("/test/webhook") + .handler(BodyHandler.create()) + .handler(ctx -> receiveWebhook(ctx)); + + router.get("/test/webhook/count") + .produces(TEXT_PLAIN) + .handler(ctx -> getWebhookCount(ctx)); } public void saveTask(String body, RoutingContext rc) { @@ -228,14 +257,86 @@ private JsonElement serialize(List spanDatas) { spanData.getResource().getAttributes().forEach((k, v) -> { jsonObject.addProperty("resource_" + k.getKey(), v.toString()); }); + + jsonObject.addProperty("events_count", spanData.getEvents().size()); + spanData.getEvents().forEach(event -> { + JsonObject eventObject = new JsonObject(); + eventObject.addProperty("name", event.getName()); + event.getAttributes().forEach((k, v) -> eventObject.addProperty("attr_" + k.getKey(), v.toString())); + jsonObject.add("event_" + event.getName(), eventObject); + }); + + jsonObject.addProperty("links_count", spanData.getLinks().size()); + if (!spanData.getLinks().isEmpty()) { + jsonObject.addProperty("link_0_spanId", spanData.getLinks().get(0).getSpanContext().getSpanId()); + jsonObject.addProperty("link_0_traceId", spanData.getLinks().get(0).getSpanContext().getTraceId()); + } + spans.add(jsonObject); } return spans; } + public void exportMetrics(RoutingContext rc) { + if (openTelemetry instanceof OpenTelemetrySdk sdk) { + sdk.getSdkMeterProvider().forceFlush(); + } + List metrics = inMemoryMetricExporter.getFinishedMetricItems().stream() + .filter(m -> m.getName().startsWith("gen_ai.")) + .collect(Collectors.toList()); + rc.response() + .setStatusCode(200) + .putHeader(CONTENT_TYPE, APPLICATION_JSON) + .end(gson.toJson(serializeMetrics(metrics))); + } + + JsonElement serializeMetrics(List metricDataList) { + JsonArray result = new JsonArray(metricDataList.size()); + for (MetricData m : metricDataList) { + JsonObject obj = new JsonObject(); + obj.addProperty("name", m.getName()); + obj.addProperty("description", m.getDescription()); + obj.addProperty("unit", m.getUnit()); + obj.addProperty("type", m.getType().name()); + if (m.getType() == MetricDataType.HISTOGRAM) { + HistogramData histogram = m.getHistogramData(); + obj.addProperty("data_count", histogram.getPoints().size()); + JsonArray points = new JsonArray(histogram.getPoints().size()); + histogram.getPoints().forEach(point -> { + JsonObject pointObj = new JsonObject(); + pointObj.addProperty("count", point.getCount()); + pointObj.addProperty("sum", point.getSum()); + point.getAttributes().forEach((k, v) -> pointObj.addProperty("attr_" + k.getKey(), v.toString())); + points.add(pointObj); + }); + obj.add("points", points); + } + result.add(obj); + } + return result; + } + + public void resetMetrics(RoutingContext rc) { + inMemoryMetricExporter.reset(); + rc.response().setStatusCode(200).end(); + } + + public void receiveWebhook(RoutingContext rc) { + webhookCallCount.incrementAndGet(); + rc.response().setStatusCode(200).end(); + } + + public void getWebhookCount(RoutingContext rc) { + rc.response() + .setStatusCode(200) + .putHeader(CONTENT_TYPE, TEXT_PLAIN) + .end(String.valueOf(webhookCallCount.get())); + } + public void reset(RoutingContext rc) { inMemorySpanExporter.reset(); + webhookCallCount.set(0); rc.response().setStatusCode(200).end(); } @@ -256,4 +357,14 @@ InMemorySpanExporter inMemorySpanExporter() { return InMemorySpanExporter.create(); } } + + @ApplicationScoped + static class InMemoryMetricExporterProducer { + + @Produces + @Singleton + InMemoryMetricExporter inMemoryMetricExporter() { + return InMemoryMetricExporter.create(); + } + } } diff --git a/extras/opentelemetry/integration-tests/src/main/java/org/a2aproject/sdk/extras/opentelemetry/it/TestAgentCardProducer.java b/extras/opentelemetry/integration-tests/src/main/java/org/a2aproject/sdk/extras/opentelemetry/it/TestAgentCardProducer.java index a34d076f0..9ea7a3b9c 100644 --- a/extras/opentelemetry/integration-tests/src/main/java/org/a2aproject/sdk/extras/opentelemetry/it/TestAgentCardProducer.java +++ b/extras/opentelemetry/integration-tests/src/main/java/org/a2aproject/sdk/extras/opentelemetry/it/TestAgentCardProducer.java @@ -31,7 +31,7 @@ public AgentCard agentCard() { .documentationUrl("http://example.com/test") .capabilities(AgentCapabilities.builder() .streaming(true) - .pushNotifications(false) + .pushNotifications(true) .build()) .defaultInputModes(Collections.singletonList("text")) .defaultOutputModes(Collections.singletonList("text")) diff --git a/extras/opentelemetry/integration-tests/src/main/resources/application.properties b/extras/opentelemetry/integration-tests/src/main/resources/application.properties index 052be4a60..aae4890aa 100644 --- a/extras/opentelemetry/integration-tests/src/main/resources/application.properties +++ b/extras/opentelemetry/integration-tests/src/main/resources/application.properties @@ -8,7 +8,8 @@ quarkus.http.test-port=8081 # OpenTelemetry configuration quarkus.otel.sdk.disabled=false quarkus.otel.traces.enabled=true -quarkus.otel.metrics.enabled=false +quarkus.otel.metrics.enabled=true +quarkus.otel.metrics.exporter=cdi quarkus.otel.logs.enabled=false quarkus.otel.instrument.vertx-http=false diff --git a/extras/opentelemetry/integration-tests/src/test/java/org/a2aproject/sdk/extras/opentelemetry/it/A2ATestRoutesTest.java b/extras/opentelemetry/integration-tests/src/test/java/org/a2aproject/sdk/extras/opentelemetry/it/A2ATestRoutesTest.java new file mode 100644 index 000000000..a0cb285ce --- /dev/null +++ b/extras/opentelemetry/integration-tests/src/test/java/org/a2aproject/sdk/extras/opentelemetry/it/A2ATestRoutesTest.java @@ -0,0 +1,96 @@ +package org.a2aproject.sdk.extras.opentelemetry.it; + +import com.google.gson.JsonArray; +import com.google.gson.JsonElement; +import com.google.gson.JsonObject; +import io.opentelemetry.api.common.Attributes; +import io.opentelemetry.sdk.metrics.data.HistogramData; +import io.opentelemetry.sdk.metrics.data.HistogramPointData; +import io.opentelemetry.sdk.metrics.data.MetricData; +import io.opentelemetry.sdk.metrics.data.MetricDataType; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.Mock; +import org.mockito.junit.jupiter.MockitoExtension; + +import java.util.List; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.Mockito.when; + +@ExtendWith(MockitoExtension.class) +class A2ATestRoutesTest { + + @Mock + private MetricData metricData; + @Mock + private HistogramData histogramData; + @Mock + private HistogramPointData point1; + @Mock + private HistogramPointData point2; + + private final A2ATestRoutes routes = new A2ATestRoutes(); + + @Test + void serializeMetrics_singleHistogramPoint_returnsPointsArray() { + when(metricData.getName()).thenReturn("gen_ai.agent.a2a.streaming.duration"); + when(metricData.getDescription()).thenReturn("desc"); + when(metricData.getUnit()).thenReturn("s"); + when(metricData.getType()).thenReturn(MetricDataType.HISTOGRAM); + when(metricData.getHistogramData()).thenReturn(histogramData); + when(histogramData.getPoints()).thenReturn(List.of(point1)); + when(point1.getCount()).thenReturn(1L); + when(point1.getSum()).thenReturn(0.5); + when(point1.getAttributes()).thenReturn(Attributes.empty()); + + JsonElement result = routes.serializeMetrics(List.of(metricData)); + + JsonObject metric = result.getAsJsonArray().get(0).getAsJsonObject(); + assertTrue(metric.has("points"), "Should have 'points' JsonArray"); + assertFalse(metric.has("point"), "Should not have single-overwriting 'point' key"); + assertEquals(1, metric.getAsJsonArray("points").size()); + } + + @Test + void serializeMetrics_multipleHistogramPoints_allPointsSerialized() { + when(metricData.getName()).thenReturn("gen_ai.agent.a2a.streaming.duration"); + when(metricData.getDescription()).thenReturn("desc"); + when(metricData.getUnit()).thenReturn("s"); + when(metricData.getType()).thenReturn(MetricDataType.HISTOGRAM); + when(metricData.getHistogramData()).thenReturn(histogramData); + when(histogramData.getPoints()).thenReturn(List.of(point1, point2)); + when(point1.getCount()).thenReturn(1L); + when(point1.getSum()).thenReturn(0.5); + when(point1.getAttributes()).thenReturn(Attributes.empty()); + when(point2.getCount()).thenReturn(3L); + when(point2.getSum()).thenReturn(2.0); + when(point2.getAttributes()).thenReturn(Attributes.empty()); + + JsonElement result = routes.serializeMetrics(List.of(metricData)); + + JsonObject metric = result.getAsJsonArray().get(0).getAsJsonObject(); + JsonArray points = metric.getAsJsonArray("points"); + assertEquals(2, points.size(), "Both histogram data points must be serialized"); + assertEquals(1L, points.get(0).getAsJsonObject().get("count").getAsLong()); + assertEquals(0.5, points.get(0).getAsJsonObject().get("sum").getAsDouble()); + assertEquals(3L, points.get(1).getAsJsonObject().get("count").getAsLong()); + assertEquals(2.0, points.get(1).getAsJsonObject().get("sum").getAsDouble()); + } + + @Test + void serializeMetrics_nonHistogramMetric_hasNoPointsKey() { + when(metricData.getName()).thenReturn("gen_ai.client.operation.duration"); + when(metricData.getDescription()).thenReturn("desc"); + when(metricData.getUnit()).thenReturn("s"); + when(metricData.getType()).thenReturn(MetricDataType.LONG_SUM); + + JsonElement result = routes.serializeMetrics(List.of(metricData)); + + JsonObject metric = result.getAsJsonArray().get(0).getAsJsonObject(); + assertFalse(metric.has("points"), "Non-histogram metric should not have a 'points' key"); + assertFalse(metric.has("point"), "Non-histogram metric should not have a 'point' key"); + } +} diff --git a/extras/opentelemetry/integration-tests/src/test/java/org/a2aproject/sdk/extras/opentelemetry/it/BaseTest.java b/extras/opentelemetry/integration-tests/src/test/java/org/a2aproject/sdk/extras/opentelemetry/it/BaseTest.java index b3c154f44..c15b66ff8 100644 --- a/extras/opentelemetry/integration-tests/src/test/java/org/a2aproject/sdk/extras/opentelemetry/it/BaseTest.java +++ b/extras/opentelemetry/integration-tests/src/test/java/org/a2aproject/sdk/extras/opentelemetry/it/BaseTest.java @@ -14,6 +14,11 @@ protected List> getSpans() { }); } + protected List> getMetrics() { + return get("/export-metrics").body().as(new TypeRef<>() { + }); + } + protected void buildGlobalTelemetryInstance() { // Do nothing in JVM mode } diff --git a/extras/opentelemetry/integration-tests/src/test/java/org/a2aproject/sdk/extras/opentelemetry/it/OpenTelemetryA2ABaseTest.java b/extras/opentelemetry/integration-tests/src/test/java/org/a2aproject/sdk/extras/opentelemetry/it/OpenTelemetryA2ABaseTest.java index cbe09435b..dd4f75d52 100644 --- a/extras/opentelemetry/integration-tests/src/test/java/org/a2aproject/sdk/extras/opentelemetry/it/OpenTelemetryA2ABaseTest.java +++ b/extras/opentelemetry/integration-tests/src/test/java/org/a2aproject/sdk/extras/opentelemetry/it/OpenTelemetryA2ABaseTest.java @@ -235,6 +235,144 @@ void testSpanAttributes() throws Exception { } } + @Test + void testSendMessageStreamCreatesClosingSpan() throws Exception { + reset(); + + Client streamingClient = Client.builder(A2A.getAgentCard("http://localhost:" + serverPort)) + .clientConfig(new ClientConfig.Builder().setStreaming(true).build()) + .withTransport(JSONRPCTransport.class, new JSONRPCTransportConfigBuilder()) + .build(); + + Message message = Message.builder() + .role(Message.Role.ROLE_USER) + .parts(List.of(new TextPart("stream test"))) + .messageId("stream-msg-1") + .build(); + MessageSendParams params = new MessageSendParams(message, null, null, ""); + + streamingClient.sendMessage(params, List.of(), null, null); + + await().atMost(10, SECONDS).until(() -> { + List> spans = getSpans(); + return spans.stream().anyMatch(s -> (A2AMethods.SEND_STREAMING_MESSAGE_METHOD + "-end").equals(s.get("name"))); + }); + + List> spans = getSpans(); + + Map initialSpan = spans.stream() + .filter(s -> A2AMethods.SEND_STREAMING_MESSAGE_METHOD.equals(s.get("name"))) + .findFirst() + .orElseThrow(() -> new AssertionError("No initial SendStreamingMessage span found")); + + Map closingSpan = spans.stream() + .filter(s -> (A2AMethods.SEND_STREAMING_MESSAGE_METHOD + "-end").equals(s.get("name"))) + .findFirst() + .orElseThrow(() -> new AssertionError("No SendStreamingMessage-stream closing span found")); + + assertEquals(Boolean.TRUE, initialSpan.get("ended"), "Initial span should be ended"); + assertEquals("SERVER", initialSpan.get("kind"), "Initial span should be SERVER"); + + assertEquals(Boolean.TRUE, closingSpan.get("ended"), "Closing span should be ended"); + assertEquals("SERVER", closingSpan.get("kind"), "Closing span should be SERVER"); + + assertTrue(((Number) closingSpan.get("events_count")).intValue() > 0, + "Closing span should have at least one streaming event"); + assertEquals(1, ((Number) closingSpan.get("links_count")).intValue(), + "Closing span should have exactly one link"); + assertEquals(initialSpan.get("spanId"), closingSpan.get("link_0_spanId"), + "Closing span should link to the initial span"); + } + + @Test + void testSendMessageStreamRecordsStreamingDuration() throws Exception { + reset(); + given().get("/reset-metrics").then().statusCode(HTTP_OK); + + Client streamingClient = Client.builder(A2A.getAgentCard("http://localhost:" + serverPort)) + .clientConfig(new ClientConfig.Builder().setStreaming(true).build()) + .withTransport(JSONRPCTransport.class, new JSONRPCTransportConfigBuilder()) + .build(); + + Message message = Message.builder() + .role(Message.Role.ROLE_USER) + .parts(List.of(new TextPart("metric stream test"))) + .messageId("metric-stream-msg-1") + .build(); + MessageSendParams params = new MessageSendParams(message, null, null, ""); + + streamingClient.sendMessage(params, List.of(), null, null); + + await().atMost(10, SECONDS).until(() -> { + List> metrics = getMetrics(); + return metrics.stream().anyMatch(m -> "gen_ai.agent.a2a.streaming.duration".equals(m.get("name"))); + }); + + List> metrics = getMetrics(); + Map streamingMetric = metrics.stream() + .filter(m -> "gen_ai.agent.a2a.streaming.duration".equals(m.get("name"))) + .findFirst() + .orElseThrow(() -> new AssertionError("No streaming duration metric found")); + + assertEquals("HISTOGRAM", streamingMetric.get("type"), + "Streaming duration metric should be a histogram"); + assertTrue(((Number) streamingMetric.get("data_count")).intValue() > 0, + "Histogram should have at least one data point"); + } + + @Test + void testPushNotificationDeliveryCreatesSpan() throws Exception { + String taskId = "push-test-task-1"; + String contextId = "push-test-ctx-1"; + + Task task = Task.builder() + .id(taskId) + .contextId(contextId) + .status(new TaskStatus(TaskState.TASK_STATE_WORKING)) + .history(Collections.emptyList()) + .artifacts(Collections.emptyList()) + .build(); + saveTaskInTaskStore(task); + ensureQueueForTask(taskId); + reset(); + + try { + TaskPushNotificationConfig config = TaskPushNotificationConfig.builder() + .id("push-test-config-1") + .taskId(taskId) + .url("http://localhost:" + serverPort + "/test/webhook") + .build(); + TaskPushNotificationConfig saved = client.createTaskPushNotificationConfiguration(config, null); + assertNotNull(saved); + + reset(); + + TaskStatusUpdateEvent event = new TaskStatusUpdateEvent( + taskId, new TaskStatus(TaskState.TASK_STATE_COMPLETED), contextId, null); + enqueueTaskStatusUpdateEventViaHttp(taskId, event); + + await().atMost(15, SECONDS).until(() -> { + List> spans = getSpans(); + return spans.stream().anyMatch(s -> "SendPushNotification".equals(s.get("name"))); + }); + + List> spans = getSpans(); + Map pushSpan = spans.stream() + .filter(s -> "SendPushNotification".equals(s.get("name"))) + .findFirst() + .orElseThrow(() -> new AssertionError("No SendPushNotification span found")); + + assertEquals(Boolean.TRUE, pushSpan.get("ended"), "Push notification span should be ended"); + assertEquals("CLIENT", pushSpan.get("kind"), "Push notification span should be CLIENT"); + assertEquals(taskId, pushSpan.get("attr_gen_ai.agent.a2a.task_id"), + "Task ID should be set on push notification span"); + assertNotNull(pushSpan.get("attr_gen_ai.agent.a2a.push_notification.event_kind"), + "Event kind should be set on push notification span"); + } finally { + deleteTaskInTaskStore(taskId); + } + } + protected void saveTaskInTaskStore(Task task) throws Exception { HttpClient httpClient = HttpClient.newBuilder() .version(HttpClient.Version.HTTP_2) @@ -297,4 +435,19 @@ protected void ensureQueueForTask(String taskId) throws Exception { throw new RuntimeException(String.format("Ensuring queue failed! Status: %d, Body: %s", response.statusCode(), response.body())); } } + + protected void enqueueTaskStatusUpdateEventViaHttp(String taskId, TaskStatusUpdateEvent event) throws Exception { + HttpClient httpClient = HttpClient.newBuilder() + .version(HttpClient.Version.HTTP_2) + .build(); + HttpRequest request = HttpRequest.newBuilder() + .uri(URI.create("http://localhost:" + serverPort + "/test/queue/enqueueTaskStatusUpdateEvent/" + taskId)) + .POST(HttpRequest.BodyPublishers.ofString(JsonUtil.toJson(event))) + .header("Content-Type", "application/json") + .build(); + HttpResponse response = httpClient.send(request, HttpResponse.BodyHandlers.ofString(StandardCharsets.UTF_8)); + if (response.statusCode() != 200) { + throw new RuntimeException(String.format("Enqueue event failed! Status: %d, Body: %s", response.statusCode(), response.body())); + } + } } diff --git a/extras/opentelemetry/integration-tests/src/test/resources/application.properties b/extras/opentelemetry/integration-tests/src/test/resources/application.properties index 932acdd58..b5eb37db9 100644 --- a/extras/opentelemetry/integration-tests/src/test/resources/application.properties +++ b/extras/opentelemetry/integration-tests/src/test/resources/application.properties @@ -7,7 +7,8 @@ quarkus.http.test-port=8081 # OpenTelemetry configuration for testing quarkus.otel.sdk.disabled=false quarkus.otel.traces.enabled=true -quarkus.otel.metrics.enabled=false +quarkus.otel.metrics.enabled=true +quarkus.otel.metrics.exporter=cdi quarkus.otel.logs.enabled=false quarkus.otel.instrument.vertx-http=false diff --git a/extras/opentelemetry/server/src/main/java/org/a2aproject/sdk/extras/opentelemetry/server/OpenTelemetryPushNotificationSenderDecorator.java b/extras/opentelemetry/server/src/main/java/org/a2aproject/sdk/extras/opentelemetry/server/OpenTelemetryPushNotificationSenderDecorator.java new file mode 100644 index 000000000..e1609dfe4 --- /dev/null +++ b/extras/opentelemetry/server/src/main/java/org/a2aproject/sdk/extras/opentelemetry/server/OpenTelemetryPushNotificationSenderDecorator.java @@ -0,0 +1,146 @@ +package org.a2aproject.sdk.extras.opentelemetry.server; + +import static org.a2aproject.sdk.extras.opentelemetry.A2AObservabilityNames.ERROR_TYPE; +import static org.a2aproject.sdk.extras.opentelemetry.A2AObservabilityNames.EXTRACT_RESPONSE_SYS_PROPERTY; +import static org.a2aproject.sdk.extras.opentelemetry.A2AObservabilityNames.GENAI_CONTEXT_ID; +import static org.a2aproject.sdk.extras.opentelemetry.A2AObservabilityNames.GENAI_OPERATION_NAME; +import static org.a2aproject.sdk.extras.opentelemetry.A2AObservabilityNames.GENAI_RESPONSE; +import static org.a2aproject.sdk.extras.opentelemetry.A2AObservabilityNames.GENAI_TASK_ID; +import static org.a2aproject.sdk.extras.opentelemetry.A2AObservabilityNames.PUSH_NOTIFICATION_EVENT_KIND; + +import io.opentelemetry.api.trace.Span; +import io.opentelemetry.api.trace.SpanKind; +import io.opentelemetry.api.trace.StatusCode; +import io.opentelemetry.api.trace.Tracer; +import io.opentelemetry.context.Scope; +import jakarta.annotation.Priority; +import jakarta.decorator.Decorator; +import jakarta.decorator.Delegate; +import jakarta.enterprise.inject.Any; +import jakarta.inject.Inject; +import org.a2aproject.sdk.server.tasks.PushNotificationSender; +import org.a2aproject.sdk.spec.Message; +import org.a2aproject.sdk.spec.StreamingEventKind; +import org.a2aproject.sdk.spec.Task; +import org.a2aproject.sdk.spec.TaskArtifactUpdateEvent; +import org.a2aproject.sdk.spec.TaskStatusUpdateEvent; +import org.jspecify.annotations.Nullable; + +/** + * OpenTelemetry CDI Decorator for {@link PushNotificationSender}. + *

+ * Creates a {@code CLIENT} span named {@code SendPushNotification} for every push + * notification delivery attempt, capturing the task ID, context ID, and event kind + * as span attributes. + *

+ * To enable this decorator, add it to your beans.xml: + * 

{@code
+ * 
+ *     org.a2aproject.sdk.extras.opentelemetry.server.OpenTelemetryPushNotificationSenderDecorator
+ * 
+ * }
+ */ +@Decorator +@Priority(100) +public abstract class OpenTelemetryPushNotificationSenderDecorator implements PushNotificationSender { + + static final String SEND_PUSH_NOTIFICATION = "SendPushNotification"; + + @Inject + @Delegate + @Any + private PushNotificationSender delegate; + + @Inject + private Tracer tracer; + + /** + * Default constructor for CDI. + */ + public OpenTelemetryPushNotificationSenderDecorator() { + } + + /** + * Constructor for unit testing. + * + * @param delegate the delegate push notification sender + * @param tracer the tracer to use + */ + OpenTelemetryPushNotificationSenderDecorator(PushNotificationSender delegate, Tracer tracer) { + this.delegate = delegate; + this.tracer = tracer; + } + + @Override + public void sendNotification(StreamingEventKind event, @Nullable Task taskSnapshot) { + var spanBuilder = tracer.spanBuilder(SEND_PUSH_NOTIFICATION) + .setSpanKind(SpanKind.CLIENT) + .setAttribute(GENAI_OPERATION_NAME, SEND_PUSH_NOTIFICATION) + .setAttribute(PUSH_NOTIFICATION_EVENT_KIND, event.getClass().getSimpleName()); + + String taskId = extractTaskId(event, taskSnapshot); + if (taskId != null) { + spanBuilder.setAttribute(GENAI_TASK_ID, taskId); + } + String contextId = extractContextId(event, taskSnapshot); + if (contextId != null) { + spanBuilder.setAttribute(GENAI_CONTEXT_ID, contextId); + } + if (extractResponse()) { + spanBuilder.setAttribute(GENAI_RESPONSE, event.toString()); + } + + Span span = spanBuilder.startSpan(); + try (Scope scope = span.makeCurrent()) { + delegate.sendNotification(event, taskSnapshot); + span.setStatus(StatusCode.OK); + } catch (Exception ex) { + span.setAttribute(ERROR_TYPE, ex.getMessage() != null ? ex.getMessage() : ex.getClass().getName()); + span.setStatus(StatusCode.ERROR, ex.getMessage()); + throw ex; + } finally { + span.end(); + } + } + + @Nullable + private static String extractTaskId(StreamingEventKind event, @Nullable Task taskSnapshot) { + if (taskSnapshot != null) { + return taskSnapshot.id(); + } + if (event instanceof Task task) { + return task.id(); + } + if (event instanceof TaskStatusUpdateEvent statusUpdate) { + return statusUpdate.taskId(); + } + if (event instanceof TaskArtifactUpdateEvent artifactUpdate) { + return artifactUpdate.taskId(); + } + return null; + } + + @Nullable + private static String extractContextId(StreamingEventKind event, @Nullable Task taskSnapshot) { + if (taskSnapshot != null) { + return taskSnapshot.contextId(); + } + if (event instanceof Task task) { + return task.contextId(); + } + if (event instanceof TaskStatusUpdateEvent statusUpdate) { + return statusUpdate.contextId(); + } + if (event instanceof TaskArtifactUpdateEvent artifactUpdate) { + return artifactUpdate.contextId(); + } + if (event instanceof Message message) { + return message.contextId(); + } + return null; + } + + private boolean extractResponse() { + return Boolean.getBoolean(EXTRACT_RESPONSE_SYS_PROPERTY); + } +} diff --git a/extras/opentelemetry/server/src/main/java/org/a2aproject/sdk/extras/opentelemetry/server/OpenTelemetryRequestHandlerDecorator.java b/extras/opentelemetry/server/src/main/java/org/a2aproject/sdk/extras/opentelemetry/server/OpenTelemetryRequestHandlerDecorator.java index 2166b06b9..f2aa94c43 100644 --- a/extras/opentelemetry/server/src/main/java/org/a2aproject/sdk/extras/opentelemetry/server/OpenTelemetryRequestHandlerDecorator.java +++ b/extras/opentelemetry/server/src/main/java/org/a2aproject/sdk/extras/opentelemetry/server/OpenTelemetryRequestHandlerDecorator.java @@ -12,10 +12,34 @@ import static org.a2aproject.sdk.extras.opentelemetry.A2AObservabilityNames.GENAI_REQUEST; import static org.a2aproject.sdk.extras.opentelemetry.A2AObservabilityNames.GENAI_RESPONSE; import static org.a2aproject.sdk.extras.opentelemetry.A2AObservabilityNames.GENAI_ROLE; +import static org.a2aproject.sdk.extras.opentelemetry.A2AObservabilityNames.GENAI_STREAMING_DURATION; +import static org.a2aproject.sdk.extras.opentelemetry.A2AObservabilityNames.GENAI_SYSTEM; +import static org.a2aproject.sdk.extras.opentelemetry.A2AObservabilityNames.GENAI_SYSTEM_VALUE; import static org.a2aproject.sdk.extras.opentelemetry.A2AObservabilityNames.GENAI_TASK_ID; +import io.opentelemetry.api.common.Attributes; +import io.opentelemetry.api.common.AttributesBuilder; +import io.opentelemetry.api.metrics.DoubleHistogram; +import io.opentelemetry.api.metrics.Meter; +import io.opentelemetry.api.trace.Span; +import io.opentelemetry.api.trace.SpanContext; +import io.opentelemetry.api.trace.SpanKind; +import io.opentelemetry.api.trace.StatusCode; +import io.opentelemetry.api.trace.Tracer; +import io.opentelemetry.context.Scope; +import jakarta.annotation.PostConstruct; +import jakarta.annotation.Priority; +import jakarta.decorator.Decorator; +import jakarta.decorator.Delegate; +import jakarta.enterprise.inject.Any; +import jakarta.enterprise.inject.Instance; +import jakarta.inject.Inject; +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.Flow; import org.a2aproject.sdk.jsonrpc.common.wrappers.ListTasksResult; import org.a2aproject.sdk.server.ServerCallContext; +import org.a2aproject.sdk.server.auth.TaskOperation; import org.a2aproject.sdk.server.requesthandlers.RequestHandler; import org.a2aproject.sdk.spec.A2AError; import org.a2aproject.sdk.spec.A2AMethods; @@ -27,23 +51,11 @@ import org.a2aproject.sdk.spec.ListTaskPushNotificationConfigsResult; import org.a2aproject.sdk.spec.ListTasksParams; import org.a2aproject.sdk.spec.MessageSendParams; -import org.a2aproject.sdk.server.auth.TaskOperation; import org.a2aproject.sdk.spec.StreamingEventKind; import org.a2aproject.sdk.spec.Task; import org.a2aproject.sdk.spec.TaskIdParams; import org.a2aproject.sdk.spec.TaskPushNotificationConfig; import org.a2aproject.sdk.spec.TaskQueryParams; -import io.opentelemetry.api.trace.Span; -import io.opentelemetry.api.trace.SpanKind; -import io.opentelemetry.api.trace.StatusCode; -import io.opentelemetry.api.trace.Tracer; -import io.opentelemetry.context.Scope; -import jakarta.annotation.Priority; -import jakarta.decorator.Decorator; -import jakarta.decorator.Delegate; -import jakarta.enterprise.inject.Any; -import jakarta.inject.Inject; -import java.util.concurrent.Flow; import org.jspecify.annotations.Nullable; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -58,6 +70,7 @@ *
  • Response data as span attributes
    • *
  • Errors and exceptions with proper status codes
    • *
  • Timing information for performance monitoring
    • + *
  • Streaming operation duration via {@code gen_ai.agent.a2a.streaming.duration} histogram
    • * *

    * To enable this decorator, add it to your beans.xml: @@ -81,6 +94,12 @@ public abstract class OpenTelemetryRequestHandlerDecorator implements RequestHan @Inject private Tracer tracer; + @Inject + private Instance meterInstance; + + @Nullable + private DoubleHistogram streamingDurationHistogram; + /** * Default constructor for CDI. */ @@ -88,14 +107,44 @@ public OpenTelemetryRequestHandlerDecorator() { } /** - * Constructor for testing. + * Constructor for testing without metrics. * * @param delegate the delegate request handler * @param tracer the tracer to use */ public OpenTelemetryRequestHandlerDecorator(RequestHandler delegate, Tracer tracer) { + this(delegate, tracer, null); + } + + /** + * Constructor for testing with metrics. + * + * @param delegate the delegate request handler + * @param tracer the tracer to use + * @param meter optional meter for streaming duration histogram; null disables metrics + */ + public OpenTelemetryRequestHandlerDecorator(RequestHandler delegate, Tracer tracer, @Nullable Meter meter) { this.delegate = delegate; this.tracer = tracer; + this.streamingDurationHistogram = buildStreamingDurationHistogram(meter); + } + + @PostConstruct + void initMetrics() { + if (streamingDurationHistogram == null && meterInstance != null && meterInstance.isResolvable()) { + streamingDurationHistogram = buildStreamingDurationHistogram(meterInstance.get()); + } + } + + @Nullable + private static DoubleHistogram buildStreamingDurationHistogram(@Nullable Meter meter) { + if (meter == null) { + return null; + } + return meter.histogramBuilder(GENAI_STREAMING_DURATION) + .setUnit("s") + .setDescription("Duration of A2A streaming operations from initiation to last event") + .build(); } @Override @@ -282,7 +331,10 @@ public Flow.Publisher onMessageSendStream(MessageSendParams } span.setStatus(StatusCode.OK); - return result; + SpanContext spanContext = span.getSpanContext(); + long startNanos = System.nanoTime(); + return new OpenTelemetryStreamPublisher(result, tracer, + A2AMethods.SEND_STREAMING_MESSAGE_METHOD, spanContext, streamingDurationHistogram, startNanos); } catch (A2AError error) { span.setAttribute(ERROR_TYPE, error.getMessage()); span.setStatus(StatusCode.ERROR, error.getMessage()); @@ -387,7 +439,10 @@ public Flow.Publisher onSubscribeToTask(TaskIdParams params, } span.setStatus(StatusCode.OK); - return result; + SpanContext spanContext = span.getSpanContext(); + long startNanos = System.nanoTime(); + return new OpenTelemetryStreamPublisher(result, tracer, + A2AMethods.SUBSCRIBE_TO_TASK_METHOD, spanContext, streamingDurationHistogram, startNanos); } catch (A2AError error) { span.setAttribute(ERROR_TYPE, error.getMessage()); span.setStatus(StatusCode.ERROR, error.getMessage()); @@ -471,4 +526,120 @@ private boolean extractRequest() { private boolean extractResponse() { return Boolean.getBoolean(EXTRACT_RESPONSE_SYS_PROPERTY); } + + private static class OpenTelemetryStreamPublisher implements Flow.Publisher { + + private final Flow.Publisher delegate; + private final Tracer tracer; + private final String spanName; + private final SpanContext parentSpanContext; + @Nullable private final DoubleHistogram streamingDurationHistogram; + private final long startNanos; + + OpenTelemetryStreamPublisher(Flow.Publisher delegate, Tracer tracer, + String spanName, SpanContext parentSpanContext, + @Nullable DoubleHistogram streamingDurationHistogram, long startNanos) { + this.delegate = delegate; + this.tracer = tracer; + this.spanName = spanName; + this.parentSpanContext = parentSpanContext; + this.streamingDurationHistogram = streamingDurationHistogram; + this.startNanos = startNanos; + } + + @Override + public void subscribe(Flow.Subscriber subscriber) { + delegate.subscribe(new OpenTelemetryStreamSubscriber(subscriber, tracer, spanName, + parentSpanContext, streamingDurationHistogram, startNanos)); + } + } + + private static class OpenTelemetryStreamSubscriber implements Flow.Subscriber { + + private final Flow.Subscriber delegate; + private final Tracer tracer; + private final String spanName; + private final SpanContext parentSpanContext; + @Nullable private final DoubleHistogram streamingDurationHistogram; + private final long startNanos; + private final List pendingEvents = new ArrayList<>(); + + OpenTelemetryStreamSubscriber(Flow.Subscriber delegate, Tracer tracer, + String spanName, SpanContext parentSpanContext, + @Nullable DoubleHistogram streamingDurationHistogram, long startNanos) { + this.delegate = delegate; + this.tracer = tracer; + this.spanName = spanName; + this.parentSpanContext = parentSpanContext; + this.streamingDurationHistogram = streamingDurationHistogram; + this.startNanos = startNanos; + } + + @Override + public void onSubscribe(Flow.Subscription subscription) { + delegate.onSubscribe(subscription); + } + + @Override + public void onNext(StreamingEventKind item) { + pendingEvents.add(new PendingEvent(spanName + "-event", + Attributes.builder() + .put("gen_ai.agent.a2a.streaming-event", item.toString()) + .put("gen_ai.agent.a2a.status.code", StatusCode.OK.name()) + .build())); + delegate.onNext(item); + } + + @Override + public void onError(Throwable throwable) { + recordStreamingDuration(false); + Span closingSpan = tracer.spanBuilder(spanName + "-end") + .setSpanKind(SpanKind.SERVER) + .addLink(parentSpanContext) + .startSpan(); + try { + for (PendingEvent event : pendingEvents) { + closingSpan.addEvent(event.name(), event.attributes()); + } + closingSpan.setStatus(StatusCode.ERROR, throwable.getMessage()); + } finally { + closingSpan.end(); + delegate.onError(throwable); + } + } + + @Override + public void onComplete() { + recordStreamingDuration(true); + Span closingSpan = tracer.spanBuilder(spanName + "-end") + .setSpanKind(SpanKind.SERVER) + .addLink(parentSpanContext) + .startSpan(); + try { + for (PendingEvent event : pendingEvents) { + closingSpan.addEvent(event.name(), event.attributes()); + } + closingSpan.setStatus(StatusCode.OK); + } finally { + closingSpan.end(); + delegate.onComplete(); + } + } + + private void recordStreamingDuration(boolean success) { + if (streamingDurationHistogram == null) { + return; + } + double seconds = (System.nanoTime() - startNanos) / 1_000_000_000.0; + AttributesBuilder builder = Attributes.builder() + .put(GENAI_OPERATION_NAME, spanName) + .put(GENAI_SYSTEM, GENAI_SYSTEM_VALUE); + if (!success) { + builder.put(ERROR_TYPE, "error"); + } + streamingDurationHistogram.record(seconds, builder.build()); + } + + private record PendingEvent(String name, Attributes attributes) {} + } } diff --git a/extras/opentelemetry/server/src/main/resources/META-INF/beans.xml b/extras/opentelemetry/server/src/main/resources/META-INF/beans.xml index e4652bc27..f775403c9 100644 --- a/extras/opentelemetry/server/src/main/resources/META-INF/beans.xml +++ b/extras/opentelemetry/server/src/main/resources/META-INF/beans.xml @@ -5,5 +5,6 @@ bean-discovery-mode="annotated"> org.a2aproject.sdk.extras.opentelemetry.server.OpenTelemetryRequestHandlerDecorator + org.a2aproject.sdk.extras.opentelemetry.server.OpenTelemetryPushNotificationSenderDecorator - \ No newline at end of file + diff --git a/extras/opentelemetry/server/src/test/java/org/a2aproject/sdk/extras/opentelemetry/server/OpenTelemetryPushNotificationSenderDecoratorTest.java b/extras/opentelemetry/server/src/test/java/org/a2aproject/sdk/extras/opentelemetry/server/OpenTelemetryPushNotificationSenderDecoratorTest.java new file mode 100644 index 000000000..3b2d8b71b --- /dev/null +++ b/extras/opentelemetry/server/src/test/java/org/a2aproject/sdk/extras/opentelemetry/server/OpenTelemetryPushNotificationSenderDecoratorTest.java @@ -0,0 +1,187 @@ +package org.a2aproject.sdk.extras.opentelemetry.server; + +import static org.a2aproject.sdk.extras.opentelemetry.A2AObservabilityNames.ERROR_TYPE; +import static org.a2aproject.sdk.extras.opentelemetry.A2AObservabilityNames.EXTRACT_RESPONSE_SYS_PROPERTY; +import static org.a2aproject.sdk.extras.opentelemetry.A2AObservabilityNames.GENAI_CONTEXT_ID; +import static org.a2aproject.sdk.extras.opentelemetry.A2AObservabilityNames.GENAI_OPERATION_NAME; +import static org.a2aproject.sdk.extras.opentelemetry.A2AObservabilityNames.GENAI_RESPONSE; +import static org.a2aproject.sdk.extras.opentelemetry.A2AObservabilityNames.GENAI_TASK_ID; +import static org.a2aproject.sdk.extras.opentelemetry.A2AObservabilityNames.PUSH_NOTIFICATION_EVENT_KIND; +import static org.a2aproject.sdk.extras.opentelemetry.server.OpenTelemetryPushNotificationSenderDecorator.SEND_PUSH_NOTIFICATION; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.lenient; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import io.opentelemetry.api.trace.Span; +import io.opentelemetry.api.trace.SpanBuilder; +import io.opentelemetry.api.trace.SpanKind; +import io.opentelemetry.api.trace.StatusCode; +import io.opentelemetry.api.trace.Tracer; +import io.opentelemetry.context.Scope; +import java.util.Collections; +import java.util.List; +import org.a2aproject.sdk.server.tasks.PushNotificationSender; +import org.a2aproject.sdk.spec.Task; +import org.a2aproject.sdk.spec.TaskArtifactUpdateEvent; +import org.a2aproject.sdk.spec.TaskState; +import org.a2aproject.sdk.spec.TaskStatus; +import org.a2aproject.sdk.spec.TaskStatusUpdateEvent; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.Mock; +import org.mockito.junit.jupiter.MockitoExtension; +import org.mockito.junit.jupiter.MockitoSettings; +import org.mockito.quality.Strictness; + +@ExtendWith(MockitoExtension.class) +@MockitoSettings(strictness = Strictness.LENIENT) +class OpenTelemetryPushNotificationSenderDecoratorTest { + + @Mock + private PushNotificationSender delegate; + + @Mock + private Tracer tracer; + + @Mock + private SpanBuilder spanBuilder; + + @Mock + private Span span; + + @Mock + private Scope scope; + + private OpenTelemetryPushNotificationSenderDecorator decorator; + + @BeforeEach + void setUp() { + lenient().when(tracer.spanBuilder(anyString())).thenReturn(spanBuilder); + lenient().when(spanBuilder.setSpanKind(any(SpanKind.class))).thenReturn(spanBuilder); + lenient().when(spanBuilder.setAttribute(anyString(), anyString())).thenReturn(spanBuilder); + lenient().when(spanBuilder.startSpan()).thenReturn(span); + lenient().when(span.makeCurrent()).thenReturn(scope); + + decorator = new TestablePushNotificationSenderDecorator(delegate, tracer); + } + + @Test + void sendNotification_createsSpanWithTaskIdFromSnapshot() { + Task taskSnapshot = Task.builder() + .id("task-1") + .contextId("ctx-1") + .status(new TaskStatus(TaskState.TASK_STATE_WORKING)) + .history(Collections.emptyList()) + .artifacts(Collections.emptyList()) + .build(); + TaskStatusUpdateEvent event = TaskStatusUpdateEvent.builder() + .taskId("task-1") + .contextId("ctx-1") + .status(new TaskStatus(TaskState.TASK_STATE_WORKING)) + .build(); + + decorator.sendNotification(event, taskSnapshot); + + verify(tracer).spanBuilder(SEND_PUSH_NOTIFICATION); + verify(spanBuilder).setSpanKind(SpanKind.CLIENT); + verify(spanBuilder).setAttribute(GENAI_OPERATION_NAME, SEND_PUSH_NOTIFICATION); + verify(spanBuilder).setAttribute(GENAI_TASK_ID, "task-1"); + verify(spanBuilder).setAttribute(GENAI_CONTEXT_ID, "ctx-1"); + verify(span).setStatus(StatusCode.OK); + verify(span).end(); + verify(delegate).sendNotification(event, taskSnapshot); + } + + @Test + void sendNotification_createsSpanWithTaskIdFromEvent() { + org.a2aproject.sdk.spec.Artifact artifact = org.a2aproject.sdk.spec.Artifact.builder() + .artifactId("artifact-id") + .parts(List.of(new org.a2aproject.sdk.spec.TextPart("result"))) + .build(); + TaskArtifactUpdateEvent event = new TaskArtifactUpdateEvent( + "task-event-1", artifact, "ctx-event-1", null, null, null); + + decorator.sendNotification(event, null); + + verify(spanBuilder).setAttribute(GENAI_TASK_ID, "task-event-1"); + verify(spanBuilder).setAttribute(GENAI_CONTEXT_ID, "ctx-event-1"); + verify(spanBuilder, never()).setAttribute(eq(GENAI_TASK_ID), eq((String) null)); + verify(span).setStatus(StatusCode.OK); + verify(span).end(); + } + + @Test + void sendNotification_setsErrorStatusOnException() { + TaskStatusUpdateEvent event = TaskStatusUpdateEvent.builder() + .taskId("task-err") + .contextId("ctx-err") + .status(new TaskStatus(TaskState.TASK_STATE_FAILED)) + .build(); + RuntimeException ex = new RuntimeException("webhook unreachable"); + doThrow(ex).when(delegate).sendNotification(any(), any()); + + assertThrows(RuntimeException.class, () -> decorator.sendNotification(event, null)); + + verify(span).setStatus(StatusCode.ERROR, "webhook unreachable"); + verify(span).setAttribute(ERROR_TYPE, "webhook unreachable"); + verify(span).end(); + } + + @Test + void sendNotification_setsEventKindAttribute() { + TaskStatusUpdateEvent event = TaskStatusUpdateEvent.builder() + .taskId("task-kind") + .contextId("ctx-kind") + .status(new TaskStatus(TaskState.TASK_STATE_COMPLETED)) + .build(); + + decorator.sendNotification(event, null); + + verify(spanBuilder).setAttribute(PUSH_NOTIFICATION_EVENT_KIND, "TaskStatusUpdateEvent"); + } + + @Test + void sendNotification_extractsResponseWhenEnabled() { + System.setProperty(EXTRACT_RESPONSE_SYS_PROPERTY, "true"); + try { + TaskStatusUpdateEvent event = TaskStatusUpdateEvent.builder() + .taskId("task-resp") + .contextId("ctx-resp") + .status(new TaskStatus(TaskState.TASK_STATE_WORKING)) + .build(); + + decorator.sendNotification(event, null); + + verify(spanBuilder).setAttribute(eq(GENAI_RESPONSE), anyString()); + } finally { + System.clearProperty(EXTRACT_RESPONSE_SYS_PROPERTY); + } + } + + @Test + void sendNotification_noResponseAttributeWhenDisabled() { + System.clearProperty(EXTRACT_RESPONSE_SYS_PROPERTY); + TaskStatusUpdateEvent event = TaskStatusUpdateEvent.builder() + .taskId("task-noresp") + .contextId("ctx-noresp") + .status(new TaskStatus(TaskState.TASK_STATE_WORKING)) + .build(); + + decorator.sendNotification(event, null); + + verify(spanBuilder, never()).setAttribute(eq(GENAI_RESPONSE), anyString()); + } + + static class TestablePushNotificationSenderDecorator extends OpenTelemetryPushNotificationSenderDecorator { + TestablePushNotificationSenderDecorator(PushNotificationSender delegate, Tracer tracer) { + super(delegate, tracer); + } + } +} diff --git a/extras/opentelemetry/server/src/test/java/org/a2aproject/sdk/extras/opentelemetry/server/OpenTelemetryRequestHandlerDecoratorTest.java b/extras/opentelemetry/server/src/test/java/org/a2aproject/sdk/extras/opentelemetry/server/OpenTelemetryRequestHandlerDecoratorTest.java index 042cd8595..c03f221f2 100644 --- a/extras/opentelemetry/server/src/test/java/org/a2aproject/sdk/extras/opentelemetry/server/OpenTelemetryRequestHandlerDecoratorTest.java +++ b/extras/opentelemetry/server/src/test/java/org/a2aproject/sdk/extras/opentelemetry/server/OpenTelemetryRequestHandlerDecoratorTest.java @@ -6,8 +6,10 @@ import static org.a2aproject.sdk.extras.opentelemetry.A2AObservabilityNames.GENAI_REQUEST; import static org.a2aproject.sdk.extras.opentelemetry.A2AObservabilityNames.GENAI_RESPONSE; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyDouble; import static org.mockito.ArgumentMatchers.anyLong; import static org.mockito.ArgumentMatchers.anyString; import static org.mockito.ArgumentMatchers.eq; @@ -20,8 +22,13 @@ import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; +import io.opentelemetry.api.common.Attributes; +import io.opentelemetry.api.metrics.DoubleHistogram; +import io.opentelemetry.api.metrics.DoubleHistogramBuilder; +import io.opentelemetry.api.metrics.Meter; import io.opentelemetry.api.trace.Span; import io.opentelemetry.api.trace.SpanBuilder; +import io.opentelemetry.api.trace.SpanContext; import io.opentelemetry.api.trace.SpanKind; import io.opentelemetry.api.trace.StatusCode; import io.opentelemetry.api.trace.Tracer; @@ -57,6 +64,7 @@ import org.junit.jupiter.api.Nested; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.ArgumentCaptor; import org.mockito.Mock; import org.mockito.junit.jupiter.MockitoExtension; @@ -75,12 +83,24 @@ class OpenTelemetryRequestHandlerDecoratorTest { @Mock private Scope scope; + @Mock + private SpanContext spanContext; + @Mock private ServerCallContext context; @Mock private RequestHandler delegate; + @Mock + private Meter meter; + + @Mock + private DoubleHistogramBuilder histogramBuilder; + + @Mock + private DoubleHistogram streamingHistogram; + private TestableOpenTelemetryRequestHandlerDecorator decorator; @BeforeEach @@ -94,14 +114,22 @@ void setUp() { lenient().when(spanBuilder.setSpanKind(any(SpanKind.class))).thenReturn(spanBuilder); lenient().when(spanBuilder.setAttribute(anyString(), anyString())).thenReturn(spanBuilder); lenient().when(spanBuilder.setAttribute(anyString(), anyLong())).thenReturn(spanBuilder); + lenient().when(spanBuilder.addLink(any(SpanContext.class))).thenReturn(spanBuilder); lenient().when(spanBuilder.startSpan()).thenReturn(span); lenient().when(span.makeCurrent()).thenReturn(scope); + lenient().when(span.getSpanContext()).thenReturn(spanContext); lenient().when(span.setAttribute(anyString(), anyString())).thenReturn(span); lenient().when(span.setStatus(any(StatusCode.class))).thenReturn(span); lenient().when(span.setStatus(any(StatusCode.class), anyString())).thenReturn(span); + // Set up meter mock chain + lenient().when(meter.histogramBuilder(anyString())).thenReturn(histogramBuilder); + lenient().when(histogramBuilder.setUnit(anyString())).thenReturn(histogramBuilder); + lenient().when(histogramBuilder.setDescription(anyString())).thenReturn(histogramBuilder); + lenient().when(histogramBuilder.build()).thenReturn(streamingHistogram); + // Create decorator with mocked dependencies - decorator = new TestableOpenTelemetryRequestHandlerDecorator(delegate, tracer); + decorator = new TestableOpenTelemetryRequestHandlerDecorator(delegate, tracer, meter); } /** @@ -111,6 +139,10 @@ static class TestableOpenTelemetryRequestHandlerDecorator extends OpenTelemetryR public TestableOpenTelemetryRequestHandlerDecorator(RequestHandler delegate, Tracer tracer) { super(delegate, tracer); } + + public TestableOpenTelemetryRequestHandlerDecorator(RequestHandler delegate, Tracer tracer, Meter meter) { + super(delegate, tracer, meter); + } } @Nested @@ -279,7 +311,7 @@ void onMessageSend_withError_setsErrorStatus() throws A2AError { @Nested class MessageSendStreamTests { @Test - void onMessageSendStream_createsSpanWithSpecialMessage() throws A2AError { + void onMessageSendStream_createsSpanAndWrapsPublisher() throws A2AError { Message message = Message.builder() .role(Message.Role.ROLE_USER) .parts(List.of(new TextPart("test message"))) @@ -293,11 +325,12 @@ void onMessageSendStream_createsSpanWithSpecialMessage() throws A2AError { Flow.Publisher actualResult = decorator.onMessageSendStream(params, context); - assertEquals(publisher, actualResult); + assertNotNull(actualResult); verify(tracer).spanBuilder(A2AMethods.SEND_STREAMING_MESSAGE_METHOD); verify(spanBuilder).setAttribute(GENAI_REQUEST, params.toString()); verify(span).setAttribute(GENAI_RESPONSE, "Stream publisher created"); verify(span).setStatus(StatusCode.OK); + verify(span).end(); } @Test @@ -319,6 +352,140 @@ void onMessageSendStream_withError_setsErrorStatus() throws A2AError { verify(span).setStatus(StatusCode.ERROR, error.getMessage()); verify(span).end(); } + + @Test + @SuppressWarnings("unchecked") + void onMessageSendStream_closingSpan_createdOnComplete() throws A2AError { + Message message = Message.builder() + .role(Message.Role.ROLE_USER) + .parts(List.of(new TextPart("test message"))) + .messageId("msg-123") + .build(); + MessageSendParams params = new MessageSendParams(message, null, null, ""); + Flow.Publisher publisher = mock(Flow.Publisher.class); + when(delegate.onMessageSendStream(params, context)).thenReturn(publisher); + + SpanBuilder closingSpanBuilder = mock(SpanBuilder.class); + Span closingSpan = mock(Span.class); + when(tracer.spanBuilder(A2AMethods.SEND_STREAMING_MESSAGE_METHOD + "-end")).thenReturn(closingSpanBuilder); + when(closingSpanBuilder.setSpanKind(any(SpanKind.class))).thenReturn(closingSpanBuilder); + when(closingSpanBuilder.addLink(any(SpanContext.class))).thenReturn(closingSpanBuilder); + when(closingSpanBuilder.startSpan()).thenReturn(closingSpan); + + Flow.Publisher result = decorator.onMessageSendStream(params, context); + + Flow.Subscriber testSubscriber = mock(Flow.Subscriber.class); + result.subscribe(testSubscriber); + + ArgumentCaptor subscriberCaptor = ArgumentCaptor.forClass(Flow.Subscriber.class); + verify(publisher).subscribe(subscriberCaptor.capture()); + Flow.Subscriber otelSubscriber = (Flow.Subscriber) subscriberCaptor.getValue(); + + otelSubscriber.onComplete(); + + verify(tracer).spanBuilder(A2AMethods.SEND_STREAMING_MESSAGE_METHOD + "-end"); + verify(closingSpanBuilder).addLink(spanContext); + verify(closingSpan).setStatus(StatusCode.OK); + verify(closingSpan).end(); + verify(testSubscriber).onComplete(); + verify(streamingHistogram).record(anyDouble(), any(Attributes.class)); + } + + @Test + @SuppressWarnings("unchecked") + void onMessageSendStream_closingSpan_recordsEventsAndEndsOnError() throws A2AError { + Message message = Message.builder() + .role(Message.Role.ROLE_USER) + .parts(List.of(new TextPart("test message"))) + .messageId("msg-123") + .build(); + MessageSendParams params = new MessageSendParams(message, null, null, ""); + Flow.Publisher publisher = mock(Flow.Publisher.class); + when(delegate.onMessageSendStream(params, context)).thenReturn(publisher); + StreamingEventKind eventItem = Message.builder() + .role(Message.Role.ROLE_AGENT) + .parts(List.of(new TextPart("event-data"))) + .messageId("event-msg-1") + .build(); + + SpanBuilder closingSpanBuilder = mock(SpanBuilder.class); + Span closingSpan = mock(Span.class); + when(tracer.spanBuilder(A2AMethods.SEND_STREAMING_MESSAGE_METHOD + "-end")).thenReturn(closingSpanBuilder); + when(closingSpanBuilder.setSpanKind(any(SpanKind.class))).thenReturn(closingSpanBuilder); + when(closingSpanBuilder.addLink(any(SpanContext.class))).thenReturn(closingSpanBuilder); + when(closingSpanBuilder.startSpan()).thenReturn(closingSpan); + + Flow.Publisher result = decorator.onMessageSendStream(params, context); + + Flow.Subscriber testSubscriber = mock(Flow.Subscriber.class); + result.subscribe(testSubscriber); + + ArgumentCaptor subscriberCaptor = ArgumentCaptor.forClass(Flow.Subscriber.class); + verify(publisher).subscribe(subscriberCaptor.capture()); + Flow.Subscriber otelSubscriber = (Flow.Subscriber) subscriberCaptor.getValue(); + + otelSubscriber.onNext(eventItem); + Throwable error = new RuntimeException("stream failure"); + otelSubscriber.onError(error); + + verify(testSubscriber).onNext(eventItem); + verify(closingSpan).addEvent(eq(A2AMethods.SEND_STREAMING_MESSAGE_METHOD + "-event"), any(Attributes.class)); + verify(closingSpan).setStatus(StatusCode.ERROR, "stream failure"); + verify(closingSpan).end(); + verify(testSubscriber).onError(error); + verify(streamingHistogram).record(anyDouble(), any(Attributes.class)); + } + + @Test + @SuppressWarnings("unchecked") + void onMessageSendStream_closingSpan_recordsMultipleEventsOnComplete() throws A2AError { + Message message = Message.builder() + .role(Message.Role.ROLE_USER) + .parts(List.of(new TextPart("test message"))) + .messageId("msg-123") + .build(); + MessageSendParams params = new MessageSendParams(message, null, null, ""); + Flow.Publisher publisher = mock(Flow.Publisher.class); + when(delegate.onMessageSendStream(params, context)).thenReturn(publisher); + StreamingEventKind event1 = Message.builder() + .role(Message.Role.ROLE_AGENT) + .parts(List.of(new TextPart("first event"))) + .messageId("event-1") + .build(); + StreamingEventKind event2 = Message.builder() + .role(Message.Role.ROLE_AGENT) + .parts(List.of(new TextPart("second event"))) + .messageId("event-2") + .build(); + + SpanBuilder closingSpanBuilder = mock(SpanBuilder.class); + Span closingSpan = mock(Span.class); + when(tracer.spanBuilder(A2AMethods.SEND_STREAMING_MESSAGE_METHOD + "-end")).thenReturn(closingSpanBuilder); + when(closingSpanBuilder.setSpanKind(any(SpanKind.class))).thenReturn(closingSpanBuilder); + when(closingSpanBuilder.addLink(any(SpanContext.class))).thenReturn(closingSpanBuilder); + when(closingSpanBuilder.startSpan()).thenReturn(closingSpan); + + Flow.Publisher result = decorator.onMessageSendStream(params, context); + + Flow.Subscriber testSubscriber = mock(Flow.Subscriber.class); + result.subscribe(testSubscriber); + + ArgumentCaptor subscriberCaptor = ArgumentCaptor.forClass(Flow.Subscriber.class); + verify(publisher).subscribe(subscriberCaptor.capture()); + Flow.Subscriber otelSubscriber = (Flow.Subscriber) subscriberCaptor.getValue(); + + otelSubscriber.onNext(event1); + otelSubscriber.onNext(event2); + otelSubscriber.onComplete(); + + verify(testSubscriber).onNext(event1); + verify(testSubscriber).onNext(event2); + verify(closingSpan, times(2)).addEvent(eq(A2AMethods.SEND_STREAMING_MESSAGE_METHOD + "-event"), any(Attributes.class)); + verify(closingSpan).setStatus(StatusCode.OK); + verify(closingSpan).end(); + verify(testSubscriber).onComplete(); + verify(streamingHistogram).record(anyDouble(), any(Attributes.class)); + } } @Nested @@ -392,18 +559,19 @@ void onGetTaskPushNotificationConfig_withError_setsErrorStatus() throws A2AError @Nested class ResubscribeToTaskTests { @Test - void onResubscribeToTask_createsSpanWithSpecialMessage() throws A2AError { + void onResubscribeToTask_createsSpanAndWrapsPublisher() throws A2AError { TaskIdParams params = new TaskIdParams("task-123"); Flow.Publisher publisher = mock(Flow.Publisher.class); when(delegate.onSubscribeToTask(params, context)).thenReturn(publisher); Flow.Publisher actualResult = decorator.onSubscribeToTask(params, context); - assertEquals(publisher, actualResult); + assertNotNull(actualResult); verify(tracer).spanBuilder(A2AMethods.SUBSCRIBE_TO_TASK_METHOD); verify(spanBuilder).setAttribute(GENAI_REQUEST, params.toString()); verify(span).setAttribute(GENAI_RESPONSE, "Stream publisher created"); verify(span).setStatus(StatusCode.OK); + verify(span).end(); } @Test @@ -418,6 +586,79 @@ void onResubscribeToTask_withError_setsErrorStatus() throws A2AError { verify(span).setStatus(StatusCode.ERROR, error.getMessage()); verify(span).end(); } + + @Test + @SuppressWarnings("unchecked") + void onResubscribeToTask_closingSpan_createdOnComplete() throws A2AError { + TaskIdParams params = new TaskIdParams("task-123"); + Flow.Publisher publisher = mock(Flow.Publisher.class); + when(delegate.onSubscribeToTask(params, context)).thenReturn(publisher); + + SpanBuilder closingSpanBuilder = mock(SpanBuilder.class); + Span closingSpan = mock(Span.class); + when(tracer.spanBuilder(A2AMethods.SUBSCRIBE_TO_TASK_METHOD + "-end")).thenReturn(closingSpanBuilder); + when(closingSpanBuilder.setSpanKind(any(SpanKind.class))).thenReturn(closingSpanBuilder); + when(closingSpanBuilder.addLink(any(SpanContext.class))).thenReturn(closingSpanBuilder); + when(closingSpanBuilder.startSpan()).thenReturn(closingSpan); + + Flow.Publisher result = decorator.onSubscribeToTask(params, context); + + Flow.Subscriber testSubscriber = mock(Flow.Subscriber.class); + result.subscribe(testSubscriber); + + ArgumentCaptor subscriberCaptor = ArgumentCaptor.forClass(Flow.Subscriber.class); + verify(publisher).subscribe(subscriberCaptor.capture()); + Flow.Subscriber otelSubscriber = (Flow.Subscriber) subscriberCaptor.getValue(); + + otelSubscriber.onComplete(); + + verify(tracer).spanBuilder(A2AMethods.SUBSCRIBE_TO_TASK_METHOD + "-end"); + verify(closingSpanBuilder).addLink(spanContext); + verify(closingSpan).setStatus(StatusCode.OK); + verify(closingSpan).end(); + verify(testSubscriber).onComplete(); + verify(streamingHistogram).record(anyDouble(), any(Attributes.class)); + } + + @Test + @SuppressWarnings("unchecked") + void onResubscribeToTask_closingSpan_recordsEventsOnCompleteAndError() throws A2AError { + TaskIdParams params = new TaskIdParams("task-123"); + Flow.Publisher publisher = mock(Flow.Publisher.class); + when(delegate.onSubscribeToTask(params, context)).thenReturn(publisher); + StreamingEventKind eventItem = Message.builder() + .role(Message.Role.ROLE_AGENT) + .parts(List.of(new TextPart("streamed-event"))) + .messageId("evt-1") + .build(); + + SpanBuilder closingSpanBuilder = mock(SpanBuilder.class); + Span closingSpan = mock(Span.class); + when(tracer.spanBuilder(A2AMethods.SUBSCRIBE_TO_TASK_METHOD + "-end")).thenReturn(closingSpanBuilder); + when(closingSpanBuilder.setSpanKind(any(SpanKind.class))).thenReturn(closingSpanBuilder); + when(closingSpanBuilder.addLink(any(SpanContext.class))).thenReturn(closingSpanBuilder); + when(closingSpanBuilder.startSpan()).thenReturn(closingSpan); + + Flow.Publisher result = decorator.onSubscribeToTask(params, context); + + Flow.Subscriber testSubscriber = mock(Flow.Subscriber.class); + result.subscribe(testSubscriber); + + ArgumentCaptor subscriberCaptor = ArgumentCaptor.forClass(Flow.Subscriber.class); + verify(publisher).subscribe(subscriberCaptor.capture()); + Flow.Subscriber otelSubscriber = (Flow.Subscriber) subscriberCaptor.getValue(); + + otelSubscriber.onNext(eventItem); + Throwable error = new RuntimeException("subscribe failure"); + otelSubscriber.onError(error); + + verify(testSubscriber).onNext(eventItem); + verify(closingSpan).addEvent(eq(A2AMethods.SUBSCRIBE_TO_TASK_METHOD + "-event"), any(Attributes.class)); + verify(closingSpan).setStatus(StatusCode.ERROR, "subscribe failure"); + verify(closingSpan).end(); + verify(testSubscriber).onError(error); + verify(streamingHistogram).record(anyDouble(), any(Attributes.class)); + } } @Nested