fix: correctly populate context_id in remote_a2a_agent library when constructing a2a request

PiperOrigin-RevId: 818813897

Author: google-genai-bot

src/google/adk/agents/remote_a2a_agent.py

@@ -337,26 +337,15 @@ def _create_a2a_request_for_user_function_response(
         ctx.session.events[-1], ctx, Role.user, self._genai_part_converter
     )
     if function_call_event.custom_metadata:
-      a2a_message.task_id = (
-          function_call_event.custom_metadata.get(
-              A2A_METADATA_PREFIX + "task_id"
-          )
-          if function_call_event.custom_metadata
-          else None
-      )
-      a2a_message.context_id = (
-          function_call_event.custom_metadata.get(
-              A2A_METADATA_PREFIX + "context_id"
-          )
-          if function_call_event.custom_metadata
-          else None
-      )
+      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")
 
     return a2a_message
 
   def _construct_message_parts_from_session(
       self, ctx: InvocationContext
-  ) -> tuple[list[A2APart], dict[str, Any], str]:
+  ) -> tuple[list[A2APart], Optional[str]]:
     """Construct A2A message parts from session events.
 
     Args:
@@ -374,11 +363,8 @@ def _construct_message_parts_from_session(
         # stop on content generated by current a2a agent given it should already
         # be in remote session
         if event.custom_metadata:
-          context_id = (
-              event.custom_metadata.get(A2A_METADATA_PREFIX + "context_id")
-              if event.custom_metadata
-              else None
-          )
+          metadata = event.custom_metadata
+          context_id = metadata.get(A2A_METADATA_PREFIX + "context_id")
         break
 
       if not event.content or not event.content.parts: