From 45d3cb37b74890be818e964a8482836b4644db77 Mon Sep 17 00:00:00 2001 From: YuqiGuo105 <131561736+YuqiGuo105@users.noreply.github.com> Date: Mon, 19 Oct 2026 23:33:12 -0600 Subject: [PATCH] feat: support metadata providers for remote A2A messages --- .../a2a/agent/A2ARequestMetadataProvider.java | 37 +++ .../google/adk/a2a/agent/RemoteA2AAgent.java | 44 +++- .../adk/a2a/agent/RemoteA2AAgentTest.java | 249 ++++++++++++++++++ .../adk/a2a/executor/AgentExecutorTest.java | 51 ++++ 4 files changed, 372 insertions(+), 9 deletions(-) create mode 100644 a2a/src/main/java/com/google/adk/a2a/agent/A2ARequestMetadataProvider.java diff --git a/a2a/src/main/java/com/google/adk/a2a/agent/A2ARequestMetadataProvider.java b/a2a/src/main/java/com/google/adk/a2a/agent/A2ARequestMetadataProvider.java new file mode 100644 index 000000000..98521bfee --- /dev/null +++ b/a2a/src/main/java/com/google/adk/a2a/agent/A2ARequestMetadataProvider.java @@ -0,0 +1,37 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.google.adk.a2a.agent; + +import com.google.adk.agents.InvocationContext; +import io.a2a.spec.Message; +import java.util.Map; +import org.jspecify.annotations.Nullable; + +/** Supplies application-selected metadata for an outgoing A2A message. */ +@FunctionalInterface +public interface A2ARequestMetadataProvider { + + /** + * Returns JSON-serializable metadata after the message's parts and identifiers have been set. + * Returning {@code null} leaves the message unchanged. An exception aborts sending the message. + * + *

The result is attached to {@link Message#getMetadata()}, not request-level metadata or + * authentication headers. On an ADK receiver it is available through the incoming A2A message; it + * is not automatically copied into session state or {@code RunConfig.customMetadata()}. + */ + @Nullable Map provide( + InvocationContext invocationContext, Message outgoingMessage); +} diff --git a/a2a/src/main/java/com/google/adk/a2a/agent/RemoteA2AAgent.java b/a2a/src/main/java/com/google/adk/a2a/agent/RemoteA2AAgent.java index bcad60b15..43c9110ed 100644 --- a/a2a/src/main/java/com/google/adk/a2a/agent/RemoteA2AAgent.java +++ b/a2a/src/main/java/com/google/adk/a2a/agent/RemoteA2AAgent.java @@ -49,7 +49,9 @@ import io.reactivex.rxjava3.core.Flowable; import io.reactivex.rxjava3.core.FlowableEmitter; import java.util.ArrayList; +import java.util.LinkedHashMap; import java.util.List; +import java.util.Map; import java.util.Objects; import java.util.Optional; import java.util.function.BiConsumer; @@ -84,6 +86,7 @@ public class RemoteA2AAgent extends BaseAgent { private final Client a2aClient; private String description; private final boolean streaming; + private final @Nullable A2ARequestMetadataProvider requestMetadataProvider; // Internal constructor used by builder private RemoteA2AAgent(Builder builder) { @@ -117,6 +120,7 @@ private RemoteA2AAgent(Builder builder) { this.description = this.agentCard.description(); } this.streaming = builder.streaming && this.agentCard.capabilities().streaming(); + this.requestMetadataProvider = builder.requestMetadataProvider; } public static Builder builder() { @@ -133,6 +137,7 @@ public static class Builder { private List beforeAgentCallback; private List afterAgentCallback; private boolean streaming; + private @Nullable A2ARequestMetadataProvider requestMetadataProvider; @CanIgnoreReturnValue public Builder streaming(boolean streaming) { @@ -182,6 +187,16 @@ public Builder a2aClient(Client a2aClient) { return this; } + /** + * Sets an optional provider called once for each outgoing message, after its parts and + * identifiers have been populated. Session state is never forwarded automatically. + */ + @CanIgnoreReturnValue + public Builder requestMetadataProvider(@Nullable A2ARequestMetadataProvider provider) { + this.requestMetadataProvider = provider; + return this; + } + public RemoteA2AAgent build() { return new RemoteA2AAgent(this); } @@ -198,19 +213,30 @@ private Message.Builder newA2AMessage( private Message prepareMessage(InvocationContext invocationContext) { Event userCall = EventConverter.findUserFunctionCall(invocationContext.session().events()); + Message message; if (userCall != null) { ImmutableList> parts = EventConverter.contentToParts(userCall.content(), userCall.partial().orElse(false)); - return newA2AMessage(invocationContext, Message.Role.USER, parts) - .taskId(EventConverter.taskId(userCall)) - .contextId(EventConverter.contextId(userCall)) - .build(); + message = + newA2AMessage(invocationContext, Message.Role.USER, parts) + .taskId(EventConverter.taskId(userCall)) + .contextId(EventConverter.contextId(userCall)) + .build(); + } else { + message = + newA2AMessage( + invocationContext, + Message.Role.USER, + EventConverter.messagePartsFromContext(invocationContext)) + .build(); + } + if (requestMetadataProvider == null) { + return message; } - return newA2AMessage( - invocationContext, - Message.Role.USER, - EventConverter.messagePartsFromContext(invocationContext)) - .build(); + Map metadata = requestMetadataProvider.provide(invocationContext, message); + return metadata == null + ? message + : new Message.Builder(message).metadata(new LinkedHashMap<>(metadata)).build(); } @Override diff --git a/a2a/src/test/java/com/google/adk/a2a/agent/RemoteA2AAgentTest.java b/a2a/src/test/java/com/google/adk/a2a/agent/RemoteA2AAgentTest.java index 2bb450003..e58b047ab 100644 --- a/a2a/src/test/java/com/google/adk/a2a/agent/RemoteA2AAgentTest.java +++ b/a2a/src/test/java/com/google/adk/a2a/agent/RemoteA2AAgentTest.java @@ -22,12 +22,15 @@ import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.verifyNoInteractions; import static org.mockito.Mockito.when; +import com.google.adk.JsonBaseModel; import com.google.adk.a2a.common.A2AClientError; import com.google.adk.a2a.common.A2AMetadata; +import com.google.adk.a2a.converters.EventConverter; import com.google.adk.agents.BaseAgent; import com.google.adk.agents.CallbackContext; import com.google.adk.agents.Callbacks; @@ -42,6 +45,7 @@ import com.google.common.collect.ImmutableMap; import com.google.genai.types.Content; import com.google.genai.types.CustomMetadata; +import com.google.genai.types.FunctionCall; import com.google.genai.types.FunctionResponse; import com.google.genai.types.Part; import io.a2a.client.Client; @@ -57,6 +61,7 @@ import io.a2a.spec.FileWithBytes; import io.a2a.spec.FileWithUri; import io.a2a.spec.Message; +import io.a2a.spec.MessageSendParams; import io.a2a.spec.Task; import io.a2a.spec.TaskArtifactUpdateEvent; import io.a2a.spec.TaskState; @@ -142,6 +147,250 @@ public void setUp() { .build(); } + @Test + public void runAsync_withoutMetadataProvider_doesNotForwardSessionState() { + session.state().put("user_id", "customer-42"); + mockStreamResponse(consumer -> consumer.accept(createFinalEvent("Done"), agentCard)); + + var unused = createAgent().runAsync(invocationContext).toList().blockingGet(); + + assertThat(captureSentMessage().getMetadata()).isNull(); + } + + @Test + public void runAsync_metadataProvider_receivesCompleteMessageAndInvocation() { + AtomicInteger calls = new AtomicInteger(); + session.state().put("user_id", "customer-42"); + session.state().put("private_key", "not-forwarded"); + RemoteA2AAgent agent = + getAgentBuilder() + .requestMetadataProvider( + (context, message) -> { + calls.incrementAndGet(); + assertThat(context.invocationId()).isEqualTo(invocationContext.invocationId()); + assertThat(context.session()).isSameInstanceAs(session); + assertThat(context.agent()).isInstanceOf(RemoteA2AAgent.class); + assertThat(message.getMessageId()).isNotEmpty(); + assertThat(message.getRole()).isEqualTo(Message.Role.USER); + assertThat(((TextPart) message.getParts().get(0)).getText()).isEqualTo("Hello"); + return Map.of("user_id", context.session().state().get("user_id")); + }) + .build(); + mockStreamResponse(consumer -> consumer.accept(createFinalEvent("Done"), agentCard)); + + var unused = agent.runAsync(invocationContext).toList().blockingGet(); + + assertThat(calls.get()).isEqualTo(1); + assertThat(captureSentMessage().getMetadata()).containsExactly("user_id", "customer-42"); + } + + @Test + public void runAsync_metadataProvider_preservesFunctionCallTaskAndContext() { + Event call = + Event.builder() + .author("remote-agent") + .content( + Content.builder() + .role("model") + .parts( + Part.builder() + .functionCall( + FunctionCall.builder().name("lookup").id("call-1").build()) + .build()) + .build()) + .customMetadata( + List.of( + CustomMetadata.builder() + .key(EventConverter.ADK_TASK_ID_KEY) + .stringValue("task-42") + .build(), + CustomMetadata.builder() + .key(EventConverter.ADK_CONTEXT_ID_KEY) + .stringValue("context-42") + .build())) + .build(); + Event response = + Event.builder() + .author("user") + .content( + Content.builder() + .role("user") + .parts( + Part.builder() + .functionResponse( + FunctionResponse.builder() + .name("lookup") + .id("call-1") + .response(Map.of("result", "found")) + .build()) + .build()) + .build()) + .build(); + Session followup = + Session.builder("followup") + .appName("demo") + .userId("user") + .events(List.of(call, response)) + .build(); + RemoteA2AAgent agent = + getAgentBuilder() + .requestMetadataProvider( + (context, message) -> { + assertThat(message.getTaskId()).isEqualTo("task-42"); + assertThat(message.getContextId()).isEqualTo("context-42"); + return Map.of("selected", true); + }) + .build(); + mockStreamResponse(consumer -> consumer.accept(createFinalEvent("Done"), agentCard)); + + var unused = + agent + .runAsync(invocationContext.toBuilder().session(followup).build()) + .toList() + .blockingGet(); + + Message sent = captureSentMessage(); + assertThat(sent.getTaskId()).isEqualTo("task-42"); + assertThat(sent.getContextId()).isEqualTo("context-42"); + assertThat(sent.getMetadata()).containsExactly("selected", true); + } + + @Test + public void runAsync_metadataProvider_isEvaluatedForEachInvocation() { + RemoteA2AAgent agent = + getAgentBuilder() + .requestMetadataProvider( + (context, message) -> Map.of("user_id", context.session().state().get("user_id"))) + .build(); + mockStreamResponse(consumer -> consumer.accept(createFinalEvent("Done"), agentCard)); + session.state().put("user_id", "first"); + var first = agent.runAsync(invocationContext).toList().blockingGet(); + session.state().put("user_id", "second"); + var second = agent.runAsync(invocationContext).toList().blockingGet(); + + ArgumentCaptor captor = ArgumentCaptor.forClass(Message.class); + verify(mockClient, times(2)) + .sendMessage(captor.capture(), any(List.class), any(Consumer.class), any()); + assertThat(captor.getAllValues().get(0).getMetadata()).containsExactly("user_id", "first"); + assertThat(captor.getAllValues().get(1).getMetadata()).containsExactly("user_id", "second"); + } + + @Test + public void runAsync_metadataProviderNullResult_leavesMetadataUnset() { + RemoteA2AAgent agent = + getAgentBuilder().requestMetadataProvider((context, message) -> null).build(); + mockStreamResponse(consumer -> consumer.accept(createFinalEvent("Done"), agentCard)); + + var unused = agent.runAsync(invocationContext).toList().blockingGet(); + + assertThat(captureSentMessage().getMetadata()).isNull(); + } + + @Test + public void runAsync_metadataProviderEmptyResult_preservesEmptyMap() { + RemoteA2AAgent agent = + getAgentBuilder().requestMetadataProvider((context, message) -> Map.of()).build(); + mockStreamResponse(consumer -> consumer.accept(createFinalEvent("Done"), agentCard)); + + var unused = agent.runAsync(invocationContext).toList().blockingGet(); + + assertThat(captureSentMessage().getMetadata()).isEmpty(); + } + + @Test + public void runAsync_metadataProviderResult_isCopiedAndAllowsJsonNull() { + Map metadata = new HashMap<>(); + metadata.put("optional", null); + metadata.put("nested", Map.of("roles", List.of("reader"))); + RemoteA2AAgent agent = + getAgentBuilder().requestMetadataProvider((context, message) -> metadata).build(); + mockStreamResponse(consumer -> consumer.accept(createFinalEvent("Done"), agentCard)); + + var unused = agent.runAsync(invocationContext).toList().blockingGet(); + metadata.put("unexpected", "later"); + + assertThat(captureSentMessage().getMetadata()) + .containsExactly("optional", null, "nested", Map.of("roles", List.of("reader"))); + } + + @Test + public void runAsync_metadataProviderFailure_abortsSend() { + IllegalStateException failure = new IllegalStateException("provider failure"); + RemoteA2AAgent agent = + getAgentBuilder() + .requestMetadataProvider( + (context, message) -> { + throw failure; + }) + .build(); + + agent.runAsync(invocationContext).test().assertError(failure); + + verify(mockClient, times(0)).sendMessage(any(), any(List.class), any(Consumer.class), any()); + } + + @Test + public void runAsync_metadataSurvivesSdkJsonRoundTripWithoutBecomingRequestMetadata() + throws Exception { + Map selected = Map.of("user_id", "customer-42", "roles", List.of("reader")); + RemoteA2AAgent agent = + getAgentBuilder().requestMetadataProvider((context, message) -> selected).build(); + mockStreamResponse(consumer -> consumer.accept(createFinalEvent("Done"), agentCard)); + var unused = agent.runAsync(invocationContext).toList().blockingGet(); + MessageSendParams params = + new MessageSendParams.Builder() + .message(captureSentMessage()) + .metadata(Map.of("request_key", "unchanged")) + .build(); + var mapper = JsonBaseModel.getMapper(); + + MessageSendParams received = + mapper.readValue(mapper.writeValueAsString(params), MessageSendParams.class); + + assertThat(received.message().getMetadata()).isEqualTo(selected); + assertThat(received.metadata()).containsExactly("request_key", "unchanged"); + assertThat(((TextPart) received.message().getParts().get(0)).getText()).isEqualTo("Hello"); + } + + @Test + public void runAsync_clearedMetadataProvider_keepsDefaultBehavior() { + RemoteA2AAgent agent = + getAgentBuilder() + .requestMetadataProvider((context, message) -> Map.of("unused", true)) + .requestMetadataProvider(null) + .build(); + mockStreamResponse(consumer -> consumer.accept(createFinalEvent("Done"), agentCard)); + + var unused = agent.runAsync(invocationContext).toList().blockingGet(); + + assertThat(captureSentMessage().getMetadata()).isNull(); + } + + @Test + public void runAsync_emptySession_doesNotEvaluateMetadataProvider() { + A2ARequestMetadataProvider provider = mock(A2ARequestMetadataProvider.class); + RemoteA2AAgent agent = getAgentBuilder().requestMetadataProvider(provider).build(); + Session empty = + Session.builder("empty").appName("demo").userId("user").events(List.of()).build(); + + assertThat( + agent + .runAsync(invocationContext.toBuilder().session(empty).build()) + .toList() + .blockingGet()) + .isEmpty(); + + verifyNoInteractions(provider); + verify(mockClient, times(0)).sendMessage(any(), any(List.class), any(Consumer.class), any()); + } + + @SuppressWarnings("unchecked") + private Message captureSentMessage() { + ArgumentCaptor captor = ArgumentCaptor.forClass(Message.class); + verify(mockClient).sendMessage(captor.capture(), any(List.class), any(Consumer.class), any()); + return captor.getValue(); + } + @Test public void createAgent_streaming_false_returnsNonStreamingAgent() { // With streaming false, the agent should not stream even if the AgentCard supports streaming. diff --git a/a2a/src/test/java/com/google/adk/a2a/executor/AgentExecutorTest.java b/a2a/src/test/java/com/google/adk/a2a/executor/AgentExecutorTest.java index 68cb196f1..8891c5680 100644 --- a/a2a/src/test/java/com/google/adk/a2a/executor/AgentExecutorTest.java +++ b/a2a/src/test/java/com/google/adk/a2a/executor/AgentExecutorTest.java @@ -26,6 +26,7 @@ import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; +import com.google.adk.JsonBaseModel; import com.google.adk.agents.BaseAgent; import com.google.adk.agents.InvocationContext; import com.google.adk.agents.RunConfig; @@ -52,8 +53,10 @@ import io.reactivex.rxjava3.core.Single; import java.util.ArrayList; import java.util.List; +import java.util.Map; import java.util.Optional; import java.util.UUID; +import java.util.concurrent.atomic.AtomicReference; import org.junit.Before; import org.junit.Test; import org.junit.runner.RunWith; @@ -408,6 +411,54 @@ public void execute_runnerSucceeds_registerCompletedTaskFails_noFailedTaskRegist assertThat(statusEvents).isEmpty(); } + @Test + public void execute_messageMetadataIsReadableAndDistinctFromRequestMetadata() throws Exception { + testAgent.setEventsToEmit(Flowable.empty()); + AtomicReference> receivedMetadata = new AtomicReference<>(); + AgentExecutor executor = + new AgentExecutor.Builder() + .agentExecutorConfig( + AgentExecutorConfig.builder() + .beforeExecuteCallback( + ctx -> { + receivedMetadata.set(ctx.getMessage().getMetadata()); + return Single.just(false); + }) + .build()) + .app(App.builder().name("test_app").rootAgent(testAgent).build()) + .sessionService(new InMemorySessionService()) + .artifactService(new InMemoryArtifactService()) + .build(); + Map messageMetadata = + Map.of("user_id", "customer-42", "roles", List.of("reader")); + Map requestMetadata = Map.of("request_key", "request-value"); + Message message = + new Message.Builder() + .messageId("msg-1") + .role(Message.Role.USER) + .parts(List.of(new TextPart("Hello"))) + .metadata(messageMetadata) + .build(); + MessageSendParams sent = + new MessageSendParams.Builder().message(message).metadata(requestMetadata).build(); + // Exercise the SDK's actual JSON models rather than flattening the two metadata maps. + var mapper = JsonBaseModel.getMapper(); + MessageSendParams received = + mapper.readValue(mapper.writeValueAsString(sent), MessageSendParams.class); + RequestContext ctx = mock(RequestContext.class); + when(ctx.getMessage()).thenReturn(received.message()); + when(ctx.getParams()).thenReturn(received); + when(ctx.getTaskId()).thenReturn("task-1"); + when(ctx.getContextId()).thenReturn("ctx-1"); + + executor.execute(ctx, eventQueue); + + assertThat(receivedMetadata.get()).isEqualTo(messageMetadata); + assertThat(testAgent.lastInvocationContext.runConfig().customMetadata()) + .containsEntry("a2a_metadata", requestMetadata); + assertThat(testAgent.lastInvocationContext.session().state()).doesNotContainKey("user_id"); + } + @Test public void execute_propagatesRequestMetadataIntoRunConfig() { testAgent.setEventsToEmit(Flowable.empty());