fix: Propagate CDI context to AgentExecutor thread · a2aproject/a2a-java@a86ea28 · GitHub
Skip to content

Commit a86ea28

Browse files
kabirkabir
committed
fix: Propagate CDI context to AgentExecutor thread
1 parent 0b867d0 commit a86ea28

11 files changed

Lines changed: 157 additions & 40 deletions

File tree

‎compat-0.3/reference/grpc/pom.xml‎

Lines changed: 4 additions & 0 deletions

‎compat-0.3/server-conversion/src/test/java/org/a2aproject/sdk/compat03/conversion/AbstractA2AServerServerTest_v0_3.java‎

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -301,6 +301,46 @@ public void testSendMessageNewMessageSuccess() throws Exception {
301301
assertEquals("test message", ((TextPart_v0_3) part).getText());
302302
}
303303

304+
@Test
305+
public void testRequestScopedBeanAvailableOnAgentExecutorThread() throws Exception {
306+
Message_v0_3 message = new Message_v0_3.Builder()
307+
.messageId("request-scoped-test")
308+
.role(Message_v0_3.Role.USER)
309+
.parts(new TextPart_v0_3("request-scoped:test"))
310+
.build();
311+
312+
CountDownLatch latch = new CountDownLatch(1);
313+
AtomicReference<Task_v0_3> receivedTask = new AtomicReference<>();
314+
AtomicReference<Throwable> errorRef = new AtomicReference<>();
315+
316+
getNonStreamingClient().sendMessage(message, List.of((event, agentCard) -> {
317+
if (event instanceof TaskEvent_v0_3 te) {
318+
receivedTask.set(te.getTask());
319+
latch.countDown();
320+
} else if (event instanceof TaskUpdateEvent_v0_3 tue) {
321+
if (tue.getUpdateEvent() instanceof TaskArtifactUpdateEvent_v0_3) {
322+
latch.countDown();
323+
}
324+
}
325+
}), error -> {
326+
errorRef.set(error);
327+
latch.countDown();
328+
});
329+
330+
assertTrue(latch.await(10, TimeUnit.SECONDS), "Request should complete within timeout");
331+
assertNull(errorRef.get(), "Should not have received an error");
332+
333+
Task_v0_3 task = receivedTask.get();
334+
assertNotNull(task, "Should have received a task");
335+
assertEquals(TaskState_v0_3.COMPLETED, task.getStatus().state());
336+
assertNotNull(task.getArtifacts());
337+
assertFalse(task.getArtifacts().isEmpty());
338+
339+
Part_v0_3<?> part = task.getArtifacts().get(0).parts().get(0);
340+
assertEquals(Part_v0_3.Kind.TEXT, part.getKind());
341+
assertEquals("request-scoped:request-scoped-value", ((TextPart_v0_3) part).getText());
342+
}
343+
304344
@Test
305345
public void testSendMessageExistingTaskSuccess() throws Exception {
306346
saveTaskInTaskStore(MINIMAL_TASK);

‎compat-0.3/server-conversion/src/test/java/org/a2aproject/sdk/compat03/conversion/AgentExecutorProducer_v0_3.java‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44

55
import jakarta.enterprise.context.ApplicationScoped;
66
import jakarta.enterprise.inject.Produces;
7+
import jakarta.inject.Inject;
78

89
import org.a2aproject.sdk.server.agentexecution.AgentExecutor;
910
import org.a2aproject.sdk.server.agentexecution.RequestContext;
@@ -17,6 +18,9 @@
1718
@IfBuildProfile("test")
1819
public class AgentExecutorProducer_v0_3 {
1920

21+
@Inject
22+
RequestScopedBean_v0_3 requestScopedBean;
23+
2024
@Produces
2125
public AgentExecutor agentExecutor() {
2226
return new AgentExecutor() {
@@ -25,6 +29,15 @@ public void execute(RequestContext context, AgentEmitter agentEmitter) throws A2
2529
String taskId = context.getTaskId();
2630
String input = context.getMessage() != null ? extractTextFromMessage(context.getMessage()) : "";
2731

32+
// Request-scoped bean test: verify CDI request context propagation
33+
if (input.startsWith("request-scoped:")) {
34+
agentEmitter.startWork();
35+
String value = requestScopedBean.getValue();
36+
agentEmitter.addArtifact(List.of(new TextPart("request-scoped:" + value)));
37+
agentEmitter.complete();
38+
return;
39+
}
40+
2841
// Special handling for multi-event test (routed by message content)
2942
if (input.startsWith("multi-event:first")) {
3043
agentEmitter.startWork();

‎compat-0.3/server-conversion/src/test/java/org/a2aproject/sdk/compat03/conversion/RequestScopedBean_v0_3.java‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,11 @@
1+
package org.a2aproject.sdk.compat03.conversion;
2+
3+
import jakarta.enterprise.context.RequestScoped;
4+
5+
@RequestScoped
6+
public class RequestScopedBean_v0_3 {
7+
8+
public String getValue() {
9+
return "request-scoped-value";
10+
}
11+
}

‎extras/opentelemetry/pom.xml‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,6 @@
1919

2020
<properties>
2121
<version.org.eclipse.microprofile.telemetry>2.0.1</version.org.eclipse.microprofile.telemetry>
22-
<version.org.eclipse.microprofile.context-propagation>1.3</version.org.eclipse.microprofile.context-propagation>
2322
</properties>
2423
<dependencies>
2524
<dependency>

‎extras/opentelemetry/server/pom.xml‎

Lines changed: 0 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -39,12 +39,6 @@
3939
<type>pom</type>
4040
<scope>provided</scope>
4141
</dependency>
42-
<dependency>
43-
<groupId>org.eclipse.microprofile.context-propagation</groupId>
44-
<artifactId>microprofile-context-propagation-api</artifactId>
45-
<version>${version.org.eclipse.microprofile.context-propagation}</version>
46-
<scope>provided</scope>
47-
</dependency>
4842
<dependency>
4943
<groupId>jakarta.enterprise</groupId>
5044
<artifactId>jakarta.enterprise.cdi-api</artifactId>

extras/opentelemetry/server/src/main/java/org/a2aproject/sdk/extras/opentelemetry/AsyncManagedExecutorProducer.java renamed to reference/common/src/main/java/org/a2aproject/sdk/server/common/quarkus/AsyncManagedExecutorProducer.java

Lines changed: 13 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -1,37 +1,28 @@
1-
package org.a2aproject.sdk.extras.opentelemetry;
1+
package org.a2aproject.sdk.server.common.quarkus;
2+
3+
import java.util.concurrent.Executor;
24

3-
import org.a2aproject.sdk.server.util.async.Internal;
45
import jakarta.annotation.PostConstruct;
56
import jakarta.annotation.Priority;
67
import jakarta.enterprise.context.ApplicationScoped;
78
import jakarta.enterprise.inject.Alternative;
89
import jakarta.enterprise.inject.Produces;
910
import jakarta.inject.Inject;
10-
import java.util.concurrent.Executor;
11+
12+
import org.a2aproject.sdk.server.util.async.Internal;
1113
import org.eclipse.microprofile.context.ManagedExecutor;
1214
import org.slf4j.Logger;
1315
import org.slf4j.LoggerFactory;
1416

1517
/**
16-
* Alternative executor producer that provides a ManagedExecutor with OpenTelemetry context propagation.
18+
* Alternative executor producer that provides a ManagedExecutor with CDI context propagation.
1719
* <p>
18-
* This producer replaces the default {@code AsyncExecutorProducer} when the OpenTelemetry extras module
19-
* is included in the application. The ManagedExecutor ensures that OpenTelemetry trace context is
20-
* properly propagated across asynchronous boundaries.
20+
* This producer replaces the default {@code AsyncExecutorProducer} so that CDI request context
21+
* (and any other registered thread context types) are propagated to the agent executor thread.
22+
* This allows {@code @RequestScoped} beans to be injected and used inside
23+
* {@link org.a2aproject.sdk.server.agentexecution.AgentExecutor#execute}.
2124
* <p>
2225
* Priority 20 ensures this alternative takes precedence over the default producer (priority 10).
23-
*
24-
* <h2>Configuration</h2>
25-
* The ManagedExecutor is container-managed and injected via CDI. Its configuration depends on the
26-
* runtime environment:
27-
* <ul>
28-
* <li><b>Quarkus:</b> Configure via {@code quarkus.thread-pool.*} properties</li>
29-
* <li><b>Other runtimes:</b> Consult your MicroProfile Context Propagation implementation documentation</li>
30-
* </ul>
31-
* <p>
32-
* Unlike the default {@code AsyncExecutorProducer}, this producer does not use the {@code a2a.executor.*}
33-
* configuration properties. The executor pool sizing and behavior are controlled by the container's
34-
* ManagedExecutor configuration.
3526
*
3627
* @see org.eclipse.microprofile.context.ManagedExecutor
3728
*/
@@ -40,13 +31,13 @@
4031
@Priority(20)
4132
public class AsyncManagedExecutorProducer {
4233
private static final Logger LOGGER = LoggerFactory.getLogger(AsyncManagedExecutorProducer.class);
43-
34+
4435
@Inject
4536
ManagedExecutor managedExecutor;
4637

4738
@PostConstruct
4839
public void init() {
49-
LOGGER.info("Initializing OpenTelemetry-aware ManagedExecutor for async operations");
40+
LOGGER.info("Initializing ManagedExecutor for async operations with CDI context propagation");
5041
if (managedExecutor == null) {
5142
LOGGER.warn("ManagedExecutor not available - context propagation may not work correctly");
5243
}
@@ -55,11 +46,10 @@ public void init() {
5546
@Produces
5647
@Internal
5748
public Executor produce() {
58-
LOGGER.debug("Using ManagedExecutor for async operations with OpenTelemetry context propagation");
49+
LOGGER.debug("Using ManagedExecutor for async operations with CDI context propagation");
5950
if (managedExecutor == null) {
6051
throw new IllegalStateException("ManagedExecutor not injected - ensure MicroProfile Context Propagation is available");
6152
}
6253
return managedExecutor;
6354
}
64-
6555
}

extras/opentelemetry/server/src/test/java/org/a2aproject/sdk/extras/opentelemetry/AsyncManagedExecutorProducerTest.java renamed to reference/common/src/test/java/org/a2aproject/sdk/server/common/quarkus/AsyncManagedExecutorProducerTest.java

Lines changed: 13 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,18 @@
1-
package org.a2aproject.sdk.extras.opentelemetry;
2-
3-
import static org.junit.jupiter.api.Assertions.*;
4-
import static org.mockito.Mockito.*;
1+
package org.a2aproject.sdk.server.common.quarkus;
2+
3+
import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
4+
import static org.junit.jupiter.api.Assertions.assertEquals;
5+
import static org.junit.jupiter.api.Assertions.assertNotNull;
6+
import static org.junit.jupiter.api.Assertions.assertNull;
7+
import static org.junit.jupiter.api.Assertions.assertSame;
8+
import static org.junit.jupiter.api.Assertions.assertThrows;
9+
import static org.junit.jupiter.api.Assertions.assertTrue;
10+
import static org.mockito.Mockito.mock;
11+
import static org.mockito.Mockito.verify;
12+
import static org.mockito.Mockito.verifyNoMoreInteractions;
513

614
import java.util.concurrent.Executor;
15+
716
import org.eclipse.microprofile.context.ManagedExecutor;
817
import org.junit.jupiter.api.BeforeEach;
918
import org.junit.jupiter.api.Nested;
@@ -39,7 +48,6 @@ void init_withValidManagedExecutor_logsSuccessfully() {
3948
void init_withNullManagedExecutor_logsWarning() {
4049
producer.managedExecutor = null;
4150

42-
// Should not throw, just log warning
4351
assertDoesNotThrow(() -> producer.init());
4452
assertNull(producer.managedExecutor);
4553
}
@@ -88,21 +96,18 @@ void produce_returnsSameInstanceOnMultipleCalls() {
8896
class CDIIntegrationTests {
8997
@Test
9098
void producer_hasCorrectAnnotations() {
91-
// Verify class has @ApplicationScoped
9299
assertTrue(
93100
AsyncManagedExecutorProducer.class.isAnnotationPresent(
94101
jakarta.enterprise.context.ApplicationScoped.class
95102
)
96103
);
97104

98-
// Verify class has @Alternative
99105
assertTrue(
100106
AsyncManagedExecutorProducer.class.isAnnotationPresent(
101107
jakarta.enterprise.inject.Alternative.class
102108
)
103109
);
104110

105-
// Verify class has @Priority with value 20
106111
assertTrue(
107112
AsyncManagedExecutorProducer.class.isAnnotationPresent(
108113
jakarta.annotation.Priority.class
@@ -120,12 +125,10 @@ void producer_hasCorrectAnnotations() {
120125
void produceMethod_hasCorrectAnnotations() throws NoSuchMethodException {
121126
var method = AsyncManagedExecutorProducer.class.getMethod("produce");
122127

123-
// Verify method has @Produces
124128
assertTrue(
125129
method.isAnnotationPresent(jakarta.enterprise.inject.Produces.class)
126130
);
127131

128-
// Verify method has @Internal
129132
assertTrue(
130133
method.isAnnotationPresent(org.a2aproject.sdk.server.util.async.Internal.class)
131134
);

‎tests/server-common/src/test/java/org/a2aproject/sdk/server/apps/common/AbstractA2AServerTest.java‎

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -504,6 +504,46 @@ public void testSendMessageNewMessageSuccess() throws Exception {
504504
assertEquals("test message", ((TextPart) part).text());
505505
}
506506

507+
@Test
508+
public void testRequestScopedBeanAvailableOnAgentExecutorThread() throws Exception {
509+
Message message = Message.builder()
510+
.messageId("request-scoped-test")
511+
.role(Message.Role.ROLE_USER)
512+
.parts(new TextPart("request-scoped:test"))
513+
.build();
514+
515+
CountDownLatch latch = new CountDownLatch(1);
516+
AtomicReference<Task> receivedTask = new AtomicReference<>();
517+
AtomicReference<Throwable> errorRef = new AtomicReference<>();
518+
519+
getNonStreamingClient().sendMessage(message, List.of((event, agentCard) -> {
520+
if (event instanceof TaskEvent te) {
521+
receivedTask.set(te.getTask());
522+
latch.countDown();
523+
} else if (event instanceof TaskUpdateEvent tue) {
524+
if (tue.getUpdateEvent() instanceof TaskArtifactUpdateEvent) {
525+
latch.countDown();
526+
}
527+
}
528+
}), error -> {
529+
errorRef.set(error);
530+
latch.countDown();
531+
});
532+
533+
assertTrue(latch.await(10, TimeUnit.SECONDS), "Request should complete within timeout");
534+
assertNull(errorRef.get(), "Should not have received an error");
535+
536+
Task task = receivedTask.get();
537+
assertNotNull(task, "Should have received a task");
538+
assertEquals(TaskState.TASK_STATE_COMPLETED, task.status().state());
539+
assertNotNull(task.artifacts());
540+
assertFalse(task.artifacts().isEmpty());
541+
542+
Part<?> part = task.artifacts().get(0).parts().get(0);
543+
assertInstanceOf(TextPart.class, part);
544+
assertEquals("request-scoped:request-scoped-value", ((TextPart) part).text());
545+
}
546+
507547
@Test
508548
public void testSendMessageExistingTaskSuccess() throws Exception {
509549
saveTaskInTaskStore(MINIMAL_TASK);

‎tests/server-common/src/test/java/org/a2aproject/sdk/server/apps/common/AgentExecutorProducer.java‎

Lines changed: 12 additions & 0 deletions

0 commit comments

Comments
 (0)