feat: Add in-process context propagation for OpenTelemetry agent execution

extras/opentelemetry/integration-tests/src/main/java/io/a2a/extras/opentelemetry/it/SimpleAgentExecutor.java

@@ -6,6 +6,7 @@
 import io.a2a.server.tasks.AgentEmitter;
 import io.a2a.spec.A2AError;
 import io.a2a.spec.TextPart;
+import io.opentelemetry.api.trace.Tracer;
 import jakarta.enterprise.context.ApplicationScoped;
 
 /**
@@ -15,6 +16,12 @@
 @ApplicationScoped
 public class SimpleAgentExecutor implements AgentExecutor {
 
+  private final Tracer tracer;
+
+  public SimpleAgentExecutor(Tracer tracer) {
+    this.tracer = tracer;
+  }
+
     @Override
     public void execute(RequestContext context, AgentEmitter emitter) throws A2AError {
         // If task doesn't exist, create it
@@ -29,6 +36,10 @@ public void execute(RequestContext context, AgentEmitter emitter) throws A2AErro
                 .findFirst()
                 .orElse("");
 
+        tracer.spanBuilder("SimpleAgentExecutor.execute")
+                .setAttribute("user.input.length", userText.length())
+                .startSpan()
+                .end();
         // Echo it back
         String response = "Echo: " + userText;
         emitter.complete(A2A.toAgentMessage(response));

extras/opentelemetry/integration-tests/src/test/java/io/a2a/extras/opentelemetry/it/OpenTelemetryA2ABaseTest.java

@@ -236,6 +236,56 @@ void testSpanAttributes() throws Exception {
         }
     }
 
+    @Test
+    void testOpenTelemetryContextIsPropagatedWhenSendMessageIsCalled() throws Exception {
+            Message testMessage = Message.builder()
+                .role(Message.Role.ROLE_USER)
+                .parts(Collections.singletonList(new TextPart("test message")))
+                .contextId("context-1234")
+                .messageId("message-1234")
+                .build();
+            client.sendMessage(testMessage);
+
+            Thread.sleep(5000);
+
+            List<Map<String, Object>> spans = getSpans();
+            var serverSpan = spans.stream().filter(span -> SpanKind.valueOf(span.get("kind").toString()) == SpanKind.SERVER).findFirst().get();
+            var agentSpan = spans.stream().filter(span -> SpanKind.valueOf(span.get("kind").toString()) == SpanKind.INTERNAL).findFirst().get();
+
+            String serverTraceId = serverSpan.get("traceId").toString();
+            String serverSpanId = serverSpan.get("spanId").toString();
+            String agentTraceId = agentSpan.get("traceId").toString();
+            String agentParentSpanId = agentSpan.get("parentSpanId").toString();
+
+            assertEquals(serverSpanId, agentParentSpanId);
+            assertEquals(serverTraceId, agentTraceId);
+    }
+
+    @Test
+    void testOpenTelemetryContextIsPropagatedWhenSendMessageStreamIsCalled() throws Exception {
+        Message testMessage = Message.builder()
+            .role(Message.Role.ROLE_USER)
+            .parts(Collections.singletonList(new TextPart("test message")))
+            .contextId("context-1234")
+            .messageId("message-1234")
+            .build();
+        client.sendMessage(testMessage);
+
+        Thread.sleep(5000);
+
+        List<Map<String, Object>> spans = getSpans();
+        var serverSpan = spans.stream().filter(span -> SpanKind.valueOf(span.get("kind").toString()) == SpanKind.SERVER).findFirst().get();
+        var agentSpan = spans.stream().filter(span -> SpanKind.valueOf(span.get("kind").toString()) == SpanKind.INTERNAL).findFirst().get();
+
+        String serverTraceId = serverSpan.get("traceId").toString();
+        String serverSpanId = serverSpan.get("spanId").toString();
+        String agentTraceId = agentSpan.get("traceId").toString();
+        String agentParentSpanId = agentSpan.get("parentSpanId").toString();
+
+        assertEquals(serverSpanId, agentParentSpanId);
+        assertEquals(serverTraceId, agentTraceId);
+    }
+
     protected void saveTaskInTaskStore(Task task) throws Exception {
         HttpClient httpClient = HttpClient.newBuilder()
                 .version(HttpClient.Version.HTTP_2)

extras/opentelemetry/server/src/main/java/io/a2a/extras/opentelemetry/OtelAgentExecutionContextPropagator.java

@@ -0,0 +1,25 @@
+package io.a2a.extras.opentelemetry;
+
+import io.a2a.server.spi.AgentExecutionContextPropagator;
+import io.opentelemetry.context.Context;
+import jakarta.annotation.Priority;
+import jakarta.enterprise.context.ApplicationScoped;
+import jakarta.enterprise.inject.Alternative;
+
+/**
+ * An implementation of AgentExecutionContextPropagator that uses OpenTelemetry's Context to propagate the execution
+ * context. This allows for proper context propagation across threads and asynchronous operations when using
+ * OpenTelemetry for tracing. The implementation simply delegates the wrapping OpenTelemetry's context propagation
+ * mechanism.
+ *
+ * @see Context#wrap(Runnable)
+ */
+@ApplicationScoped
+@Alternative
+@Priority(1)
+public class OtelAgentExecutionContextPropagator implements AgentExecutionContextPropagator {
+  @Override
+  public Runnable wrap(Runnable runnable) {
+    return Context.current().wrap(runnable);
+  }
+}

extras/opentelemetry/server/src/test/java/io/a2a/extras/opentelemetry/OtelAgentExecutionContextPropagatorTest.java

@@ -0,0 +1,46 @@
+package io.a2a.extras.opentelemetry;
+
+import io.opentelemetry.context.Context;
+import io.opentelemetry.context.ContextKey;
+import io.opentelemetry.context.Scope;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+class OtelAgentExecutionContextPropagatorTest {
+
+  ContextKey<String> exampleContextKey = ContextKey.named("test-key");
+  Scope scope = null;
+
+  @BeforeEach
+  void setup() {
+    Context context = Context.current().with(exampleContextKey, "test-value");
+    scope = context.makeCurrent();
+  }
+
+  @Test
+  void testOtelAgentExecutionContextPropagation() throws InterruptedException {
+
+    OtelAgentExecutionContextPropagator propagator = new OtelAgentExecutionContextPropagator();
+    Runnable wrappedContext =
+        propagator.wrap(() -> assertEquals("test-value", Context.current().get(exampleContextKey)));
+
+    ExecutorService executorService = Executors.newSingleThreadExecutor();
+    CompletableFuture<Void> future = CompletableFuture.runAsync(wrappedContext, executorService);
+    future.join();
+    executorService.shutdown();
+    executorService.awaitTermination(5, TimeUnit.SECONDS);
+  }
+
+  @AfterEach
+  void tearDown() {
+    scope.close();
+  }
+}

server-common/src/main/java/io/a2a/server/context/DefaultAgentExecutionContextPropagator.java

@@ -0,0 +1,13 @@
+package io.a2a.server.context;
+
+import io.a2a.server.spi.AgentExecutionContextPropagator;
+import jakarta.enterprise.context.ApplicationScoped;
+
+@ApplicationScoped
+public class DefaultAgentExecutionContextPropagator implements AgentExecutionContextPropagator {
+
+  @Override
+  public Runnable wrap(Runnable runnable) {
+    return runnable;
+  }
+}

server-common/src/main/java/io/a2a/server/requesthandlers/DefaultRequestHandler.java

@@ -20,6 +20,8 @@
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Supplier;
 
+import io.a2a.server.context.DefaultAgentExecutionContextPropagator;
+import io.a2a.server.spi.AgentExecutionContextPropagator;
 import jakarta.annotation.PostConstruct;
 import jakarta.enterprise.context.ApplicationScoped;
 import jakarta.inject.Inject;
@@ -225,6 +227,7 @@ public class DefaultRequestHandler implements RequestHandler {
     private PushNotificationConfigStore pushConfigStore;
     private MainEventBusProcessor mainEventBusProcessor;
     private Supplier<RequestContext.Builder> requestContextBuilder;
+    private AgentExecutionContextPropagator context;
 
     private final ConcurrentMap<String, CompletableFuture<Void>> runningAgents = new ConcurrentHashMap<>();
 
@@ -248,14 +251,16 @@ protected DefaultRequestHandler() {
         this.requestContextBuilder = null;
         this.executor = null;
         this.eventConsumerExecutor = null;
+        this.context = null;
     }
 
     @Inject
     public DefaultRequestHandler(AgentExecutor agentExecutor, TaskStore taskStore,
                                  QueueManager queueManager, PushNotificationConfigStore pushConfigStore,
                                  MainEventBusProcessor mainEventBusProcessor,
                                  @Internal Executor executor,
-                                 @EventConsumerExecutor Executor eventConsumerExecutor) {
+                                 @EventConsumerExecutor Executor eventConsumerExecutor,
+                                 AgentExecutionContextPropagator context) {
         this.agentExecutor = agentExecutor;
         this.taskStore = taskStore;
         this.queueManager = queueManager;
@@ -268,6 +273,7 @@ public DefaultRequestHandler(AgentExecutor agentExecutor, TaskStore taskStore,
         //  I am unsure about the correct scope.
         //  Also reworked to make a Supplier since otherwise the builder gets polluted with wrong tasks
         this.requestContextBuilder = () -> new SimpleRequestContextBuilder(taskStore, false);
+        this.context = context;
     }
 
     @PostConstruct
@@ -288,7 +294,7 @@ public static DefaultRequestHandler create(AgentExecutor agentExecutor, TaskStor
                          Executor executor, Executor eventConsumerExecutor) {
         DefaultRequestHandler handler =
                 new DefaultRequestHandler(agentExecutor, taskStore, queueManager, pushConfigStore,
-                        mainEventBusProcessor, executor, eventConsumerExecutor);
+                        mainEventBusProcessor, executor, eventConsumerExecutor, new DefaultAgentExecutionContextPropagator());
         handler.agentCompletionTimeoutSeconds = 5;
         handler.consumptionCompletionTimeoutSeconds = 2;
 
@@ -883,9 +889,8 @@ private boolean shouldAddPushInfo(MessageSendParams params) {
     private EnhancedRunnable registerAndExecuteAgentAsync(String taskId, RequestContext requestContext, EventQueue queue, EnhancedRunnable.DoneCallback doneCallback) {
         LOGGER.debug("Registering agent execution for task {}, runningAgents.size() before: {}", taskId, runningAgents.size());
         logThreadStats("AGENT START");
-        EnhancedRunnable runnable = new EnhancedRunnable() {
-            @Override
-            public void run() {
+
+        Runnable wrappedRunnable = context.wrap(() -> {
                 LOGGER.debug("Agent execution starting for task {}", taskId);
                 AgentEmitter emitter = new AgentEmitter(requestContext, queue);
                 try {
@@ -909,6 +914,11 @@ public void run() {
                 LOGGER.debug("Agent execution completed for task {}", taskId);
                 // The consumer (running on the Vert.x worker thread) handles queue lifecycle.
                 // This avoids blocking agent-executor threads waiting for worker threads.
+        });
+        EnhancedRunnable runnable = new EnhancedRunnable() {
+            @Override
+            public void run() {
+                wrappedRunnable.run();
             }
         };
 

server-common/src/main/java/io/a2a/server/spi/AgentExecutionContextPropagator.java

@@ -0,0 +1,14 @@
+package io.a2a.server.spi;
+
+/**
+ * Interface for propagating the agent execution context across different threads or asynchronous operations.
+ * Implementations can wrap Runnable tasks to ensure that the context is properly propagated.
+ */
+public interface AgentExecutionContextPropagator {
+
+  /**
+   * Wraps a Runnable task to propagate the agent execution context when the task is executed.
+   * @param runnable the Runnable task to wrap
+   */
+  Runnable wrap(Runnable runnable);
+}

server-common/src/test/java/io/a2a/server/context/DefaultAgentExecutionContextPropagatorTest.java

@@ -0,0 +1,15 @@
+package io.a2a.server.context;
+
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+
+class DefaultAgentExecutionContextPropagatorTest {
+
+  @Test
+  void testDefaultAgentExecutionPropagation() {
+    DefaultAgentExecutionContextPropagator propagator = new DefaultAgentExecutionContextPropagator();
+    Runnable runnable = () -> {
+    };
+    Assertions.assertSame(runnable, propagator.wrap(runnable));
+  }
+}