Skip to content

Commit 02c294a

Browse files
ehsavoieehsavoie
committed
WIP
Signed-off-by: Emmanuel Hugonnet <[email protected]>
1 parent 6e36a9c commit 02c294a

File tree

4 files changed

+56
-11
lines changed

4 files changed

+56
-11
lines changed

‎transport/grpc/src/main/java/io/a2a/transport/grpc/handler/GrpcHandler.java‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
import static io.a2a.grpc.utils.ProtoUtils.FromProto;
44
import static io.a2a.grpc.utils.ProtoUtils.ToProto;
5+
import static io.a2a.server.interceptors.Kind.SERVER;
56

67
import jakarta.enterprise.inject.Vetoed;
78

@@ -27,6 +28,7 @@
2728
import io.a2a.server.auth.UnauthenticatedUser;
2829
import io.a2a.server.auth.User;
2930
import io.a2a.server.extensions.A2AExtensions;
31+
import io.a2a.server.interceptors.Trace;
3032
import io.a2a.transport.grpc.context.GrpcContextKeys;
3133
import io.a2a.server.requesthandlers.RequestHandler;
3234
import io.a2a.spec.AgentCard;
@@ -71,6 +73,7 @@ public GrpcHandler() {
7173
}
7274

7375
@Override
76+
@Trace(extractor = GrpctAttributeExtractor.class, kind = SERVER)
7477
public void sendMessage(io.a2a.grpc.SendMessageRequest request,
7578
StreamObserver<io.a2a.grpc.SendMessageResponse> responseObserver) {
7679
try {
@@ -90,6 +93,7 @@ public void sendMessage(io.a2a.grpc.SendMessageRequest request,
9093
}
9194

9295
@Override
96+
@Trace(extractor = GrpctAttributeExtractor.class, kind = SERVER)
9397
public void getTask(io.a2a.grpc.GetTaskRequest request,
9498
StreamObserver<io.a2a.grpc.Task> responseObserver) {
9599
try {
@@ -112,6 +116,7 @@ public void getTask(io.a2a.grpc.GetTaskRequest request,
112116
}
113117

114118
@Override
119+
@Trace(extractor = GrpctAttributeExtractor.class, kind = SERVER)
115120
public void listTasks(io.a2a.grpc.ListTasksRequest request,
116121
StreamObserver<io.a2a.grpc.ListTasksResponse> responseObserver) {
117122
try {
@@ -130,6 +135,7 @@ public void listTasks(io.a2a.grpc.ListTasksRequest request,
130135
}
131136

132137
@Override
138+
@Trace(extractor = GrpctAttributeExtractor.class, kind = SERVER)
133139
public void cancelTask(io.a2a.grpc.CancelTaskRequest request,
134140
StreamObserver<io.a2a.grpc.Task> responseObserver) {
135141
try {
@@ -152,6 +158,7 @@ public void cancelTask(io.a2a.grpc.CancelTaskRequest request,
152158
}
153159

154160
@Override
161+
@Trace(extractor = GrpctAttributeExtractor.class, kind = SERVER)
155162
public void createTaskPushNotificationConfig(io.a2a.grpc.CreateTaskPushNotificationConfigRequest request,
156163
StreamObserver<io.a2a.grpc.TaskPushNotificationConfig> responseObserver) {
157164
if (!getAgentCardInternal().capabilities().pushNotifications()) {
@@ -175,6 +182,7 @@ public void createTaskPushNotificationConfig(io.a2a.grpc.CreateTaskPushNotificat
175182
}
176183

177184
@Override
185+
@Trace(extractor = GrpctAttributeExtractor.class, kind = SERVER)
178186
public void getTaskPushNotificationConfig(io.a2a.grpc.GetTaskPushNotificationConfigRequest request,
179187
StreamObserver<io.a2a.grpc.TaskPushNotificationConfig> responseObserver) {
180188
if (!getAgentCardInternal().capabilities().pushNotifications()) {
@@ -198,6 +206,7 @@ public void getTaskPushNotificationConfig(io.a2a.grpc.GetTaskPushNotificationCon
198206
}
199207

200208
@Override
209+
@Trace(extractor = GrpctAttributeExtractor.class, kind = SERVER)
201210
public void listTaskPushNotificationConfig(io.a2a.grpc.ListTaskPushNotificationConfigRequest request,
202211
StreamObserver<io.a2a.grpc.ListTaskPushNotificationConfigResponse> responseObserver) {
203212
if (!getAgentCardInternal().capabilities().pushNotifications()) {
@@ -226,6 +235,7 @@ public void listTaskPushNotificationConfig(io.a2a.grpc.ListTaskPushNotificationC
226235
}
227236

228237
@Override
238+
@Trace(extractor = GrpctAttributeExtractor.class, kind = SERVER)
229239
public void sendStreamingMessage(io.a2a.grpc.SendMessageRequest request,
230240
StreamObserver<io.a2a.grpc.StreamResponse> responseObserver) {
231241
if (!getAgentCardInternal().capabilities().streaming()) {
@@ -248,6 +258,7 @@ public void sendStreamingMessage(io.a2a.grpc.SendMessageRequest request,
248258
}
249259

250260
@Override
261+
@Trace(extractor = GrpctAttributeExtractor.class, kind = SERVER)
251262
public void taskSubscription(io.a2a.grpc.TaskSubscriptionRequest request,
252263
StreamObserver<io.a2a.grpc.StreamResponse> responseObserver) {
253264
if (!getAgentCardInternal().capabilities().streaming()) {
@@ -328,6 +339,7 @@ public void getAgentCard(io.a2a.grpc.GetAgentCardRequest request,
328339
}
329340

330341
@Override
342+
@Trace(extractor = GrpctAttributeExtractor.class, kind = SERVER)
331343
public void deleteTaskPushNotificationConfig(io.a2a.grpc.DeleteTaskPushNotificationConfigRequest request,
332344
StreamObserver<Empty> responseObserver) {
333345
if (!getAgentCardInternal().capabilities().pushNotifications()) {

‎transport/grpc/src/main/java/io/a2a/transport/grpc/handler/GrpctAttributeExtractor.java‎

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,33 @@
1+
package io.a2a.transport.grpc.handler;
2+
3+
import io.a2a.server.interceptors.AttributeExtractor;
4+
import io.grpc.Context;
5+
import java.util.Collections;
6+
import java.util.Map;
7+
import io.a2a.transport.grpc.context.GrpcContextKeys;
8+
9+
public class GrpctAttributeExtractor implements AttributeExtractor {
10+
11+
@Override
12+
public Map<String, String> extract(Object target, String method, Object[] parameters) {
13+
switch (method) {
14+
case "sendMessage",
15+
"getTask",
16+
"listTasks",
17+
"cancelTask",
18+
"createTaskPushNotificationConfig",
19+
"getTaskPushNotificationConfig",
20+
"listTaskPushNotificationConfig",
21+
"sendStreamingMessage",
22+
"taskSubscription",
23+
"deleteTaskPushNotificationConfig" -> {
24+
Context currentContext = Context.current();
25+
return Map.of("body", parameters[0].toString(), "extensions", GrpcContextKeys.EXTENSIONS_HEADER_KEY.get(), "a2a.method", GrpcContextKeys.METHOD_NAME_KEY.get(currentContext));
26+
}
27+
28+
default -> {
29+
return Collections.emptyMap();
30+
}
31+
}
32+
}
33+
}

‎transport/jsonrpc/src/main/java/io/a2a/transport/jsonrpc/handler/JSONRPCHandler.java‎

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -79,7 +79,7 @@ public JSONRPCHandler(@PublicAgentCard AgentCard agentCard, RequestHandler reque
7979
this(agentCard, null, requestHandler, executor);
8080
}
8181

82-
@Trace(extractor=RestAttributeExtractor.class, kind = SERVER)
82+
@Trace(extractor=JsonRPCAttributeExtractor.class, kind = SERVER)
8383
public SendMessageResponse onMessageSend(SendMessageRequest request, ServerCallContext context) {
8484
try {
8585
EventKind taskOrMessage = requestHandler.onMessageSend(request.getParams(), context);
@@ -91,7 +91,7 @@ public SendMessageResponse onMessageSend(SendMessageRequest request, ServerCallC
9191
}
9292
}
9393

94-
@Trace(extractor=RestAttributeExtractor.class, kind = SERVER)
94+
@Trace(extractor=JsonRPCAttributeExtractor.class, kind = SERVER)
9595
public Flow.Publisher<SendStreamingMessageResponse> onMessageSendStream(
9696
SendStreamingMessageRequest request, ServerCallContext context) {
9797
if (!agentCard.capabilities().streaming()) {
@@ -114,7 +114,7 @@ public Flow.Publisher<SendStreamingMessageResponse> onMessageSendStream(
114114
}
115115
}
116116

117-
@Trace(extractor=RestAttributeExtractor.class, kind = SERVER)
117+
@Trace(extractor=JsonRPCAttributeExtractor.class, kind = SERVER)
118118
public CancelTaskResponse onCancelTask(CancelTaskRequest request, ServerCallContext context) {
119119
try {
120120
Task task = requestHandler.onCancelTask(request.getParams(), context);
@@ -129,7 +129,7 @@ public CancelTaskResponse onCancelTask(CancelTaskRequest request, ServerCallCont
129129
}
130130
}
131131

132-
@Trace(extractor=RestAttributeExtractor.class, kind = SERVER)
132+
@Trace(extractor=JsonRPCAttributeExtractor.class, kind = SERVER)
133133
public Flow.Publisher<SendStreamingMessageResponse> onResubscribeToTask(
134134
TaskResubscriptionRequest request, ServerCallContext context) {
135135
if (!agentCard.capabilities().streaming()) {
@@ -151,7 +151,7 @@ public Flow.Publisher<SendStreamingMessageResponse> onResubscribeToTask(
151151
}
152152
}
153153

154-
@Trace(extractor=RestAttributeExtractor.class, kind = SERVER)
154+
@Trace(extractor=JsonRPCAttributeExtractor.class, kind = SERVER)
155155
public GetTaskPushNotificationConfigResponse getPushNotificationConfig(
156156
GetTaskPushNotificationConfigRequest request, ServerCallContext context) {
157157
if (!agentCard.capabilities().pushNotifications()) {
@@ -169,7 +169,7 @@ public GetTaskPushNotificationConfigResponse getPushNotificationConfig(
169169
}
170170
}
171171

172-
@Trace(extractor=RestAttributeExtractor.class, kind = SERVER)
172+
@Trace(extractor=JsonRPCAttributeExtractor.class, kind = SERVER)
173173
public SetTaskPushNotificationConfigResponse setPushNotificationConfig(
174174
SetTaskPushNotificationConfigRequest request, ServerCallContext context) {
175175
if (!agentCard.capabilities().pushNotifications()) {
@@ -187,7 +187,7 @@ public SetTaskPushNotificationConfigResponse setPushNotificationConfig(
187187
}
188188
}
189189

190-
@Trace(extractor=RestAttributeExtractor.class, kind = SERVER)
190+
@Trace(extractor=JsonRPCAttributeExtractor.class, kind = SERVER)
191191
public GetTaskResponse onGetTask(GetTaskRequest request, ServerCallContext context) {
192192
try {
193193
Task task = requestHandler.onGetTask(request.getParams(), context);
@@ -199,7 +199,7 @@ public GetTaskResponse onGetTask(GetTaskRequest request, ServerCallContext conte
199199
}
200200
}
201201

202-
@Trace(extractor=RestAttributeExtractor.class, kind = SERVER)
202+
@Trace(extractor=JsonRPCAttributeExtractor.class, kind = SERVER)
203203
public ListTasksResponse onListTasks(ListTasksRequest request, ServerCallContext context) {
204204
try {
205205
ListTasksResult result = requestHandler.onListTasks(request.getParams(), context);
@@ -211,7 +211,7 @@ public ListTasksResponse onListTasks(ListTasksRequest request, ServerCallContext
211211
}
212212
}
213213

214-
@Trace(extractor=RestAttributeExtractor.class, kind = SERVER)
214+
@Trace(extractor=JsonRPCAttributeExtractor.class, kind = SERVER)
215215
public ListTaskPushNotificationConfigResponse listPushNotificationConfig(
216216
ListTaskPushNotificationConfigRequest request, ServerCallContext context) {
217217
if ( !agentCard.capabilities().pushNotifications()) {
@@ -229,7 +229,7 @@ public ListTaskPushNotificationConfigResponse listPushNotificationConfig(
229229
}
230230
}
231231

232-
@Trace(extractor=RestAttributeExtractor.class, kind = SERVER)
232+
@Trace(extractor=JsonRPCAttributeExtractor.class, kind = SERVER)
233233
public DeleteTaskPushNotificationConfigResponse deletePushNotificationConfig(
234234
DeleteTaskPushNotificationConfigRequest request, ServerCallContext context) {
235235
if ( !agentCard.capabilities().pushNotifications()) {

transport/jsonrpc/src/main/java/io/a2a/transport/jsonrpc/handler/RestAttributeExtractor.java renamed to transport/jsonrpc/src/main/java/io/a2a/transport/jsonrpc/handler/JsonRPCAttributeExtractor.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@
88
import java.util.Map;
99
import java.util.stream.Collectors;
1010

11-
public class RestAttributeExtractor implements AttributeExtractor {
11+
public class JsonRPCAttributeExtractor implements AttributeExtractor {
1212

1313
@Override
1414
public Map<String, String> extract(Object target, String method, Object[] parameters) {

0 commit comments

Comments
 (0)