feat: Add support for passing request metadata in RemoteA2AAgent

This change updates `RemoteA2AAgent` to extract and forward custom metadata from session events to the `a2a-sdk`'s `send_message` method. The metadata is looked for under the key `A2A_METADATA_PREFIX + "metadata"` within the `custom_metadata` of the relevant session events. The `a2a-sdk` dependency is also updated to a version that supports this feature.

This feature was added in v0.3.11 of the a2a-sdk library: https://github.com/a2aproject/a2a-python/releases/tag/v0.3.11

PiperOrigin-RevId: 831120978
This commit is contained in:
Google Team Member
2025-11-11 16:49:39 -08:00
committed by Copybara-Service
parent 249216e890
commit d85a301039
3 changed files with 84 additions and 57 deletions
+19 -11
View File
@@ -318,7 +318,7 @@ class RemoteA2aAgent(BaseAgent):
def _create_a2a_request_for_user_function_response(
self, ctx: InvocationContext
) -> Optional[A2AMessage]:
) -> tuple[Optional[A2AMessage], Optional[dict[str, Any]]]:
"""Create A2A request for user function response if applicable.
Args:
@@ -328,34 +328,38 @@ class RemoteA2aAgent(BaseAgent):
SendMessageRequest if function response found, None otherwise
"""
if not ctx.session.events or ctx.session.events[-1].author != "user":
return None
return None, None
function_call_event = find_matching_function_call(ctx.session.events)
if not function_call_event:
return None
return None, None
a2a_message = convert_event_to_a2a_message(
ctx.session.events[-1], ctx, Role.user, self._genai_part_converter
)
message_metadata = None
if function_call_event.custom_metadata:
metadata = function_call_event.custom_metadata
a2a_message.task_id = metadata.get(A2A_METADATA_PREFIX + "task_id")
a2a_message.context_id = metadata.get(A2A_METADATA_PREFIX + "context_id")
message_metadata = metadata.get(A2A_METADATA_PREFIX + "metadata")
return a2a_message
return a2a_message, message_metadata
def _construct_message_parts_from_session(
self, ctx: InvocationContext
) -> tuple[list[A2APart], Optional[str]]:
) -> tuple[list[A2APart], Optional[str], Optional[dict[str, Any]]]:
"""Construct A2A message parts from session events.
Args:
ctx: The invocation context
Returns:
List of A2A parts extracted from session events, context ID
List of A2A parts extracted from session events, context ID,
request metadata
"""
message_parts: list[A2APart] = []
context_id = None
request_metadata = None
events_to_process = []
for event in reversed(ctx.session.events):
@@ -365,6 +369,7 @@ class RemoteA2aAgent(BaseAgent):
if event.custom_metadata:
metadata = event.custom_metadata
context_id = metadata.get(A2A_METADATA_PREFIX + "context_id")
request_metadata = metadata.get(A2A_METADATA_PREFIX + "metadata")
break
events_to_process.append(event)
@@ -385,7 +390,7 @@ class RemoteA2aAgent(BaseAgent):
else:
logger.warning("Failed to convert part to A2A format: %s", part)
return message_parts, context_id
return message_parts, context_id, request_metadata
async def _handle_a2a_response(
self, a2a_response: A2AClientEvent | A2AMessage, ctx: InvocationContext
@@ -493,10 +498,12 @@ class RemoteA2aAgent(BaseAgent):
return
# Create A2A request for function response or regular message
a2a_request = self._create_a2a_request_for_user_function_response(ctx)
a2a_request, request_metadata = (
self._create_a2a_request_for_user_function_response(ctx)
)
if not a2a_request:
message_parts, context_id = self._construct_message_parts_from_session(
ctx
message_parts, context_id, request_metadata = (
self._construct_message_parts_from_session(ctx)
)
if not message_parts:
@@ -522,7 +529,8 @@ class RemoteA2aAgent(BaseAgent):
try:
async for a2a_response in self._a2a_client.send_message(
request=a2a_request
request=a2a_request,
request_metadata=request_metadata,
):
logger.debug(build_a2a_response_log(a2a_response))