feat(a2a): Add optional context_builder to RemoteA2aAgent

src/google/adk/agents/remote_a2a_agent.py

@@ -22,6 +22,7 @@
 from typing import AsyncGenerator
 from typing import Callable
 from typing import Optional
+from typing import Protocol
 from typing import Union
 from urllib.parse import urlparse
 
@@ -77,6 +78,7 @@
     "A2AClientError",
     "AGENT_CARD_WELL_KNOWN_PATH",
     "AgentCardResolutionError",
+    "ContextBuilder",
     "RemoteA2aAgent",
 ]
 
@@ -137,6 +139,35 @@ class A2AClientError(Exception):
   pass
 
 
+class ContextBuilder(Protocol):
+  """Protocol for custom A2A request context builders on ``RemoteA2aAgent``.
+
+  When provided to ``RemoteA2aAgent``, replaces the default session→message
+  construction in ``_construct_message_parts_from_session``. Use this to send
+  only the current turn, a sliding window, a summarized history, or any other
+  domain-specific context strategy.
+  """
+
+  def __call__(
+      self,
+      ctx: InvocationContext,
+      agent_name: str,
+      genai_part_converter: GenAIPartToA2APartConverter,
+  ) -> tuple[list[A2APart], Optional[str]]:
+    """Build A2A message parts from the ADK session context.
+
+    Args:
+      ctx: The invocation context containing session events.
+      agent_name: Name of the current ``RemoteA2aAgent`` instance.
+      genai_part_converter: Function to convert GenAI parts to A2A parts.
+
+    Returns:
+      Tuple of ``(message_parts, context_id)``. ``context_id`` may be ``None``
+      for a new / stateless remote session.
+    """
+    ...
+
+
 def _add_mock_function_call(event: Event, state: TaskState) -> None:
   """Generates a mock function call for input-required events if applicable."""
   if event.content is None:
@@ -184,6 +215,7 @@ def __init__(
           Callable[[InvocationContext, A2AMessage], dict[str, Any]]
       ] = None,
       full_history_when_stateless: bool = False,
+      context_builder: Optional[ContextBuilder] = None,
       config: Optional[A2aRemoteAgentConfig] = None,
       use_legacy: bool = True,
       **kwargs: Any,
@@ -206,6 +238,10 @@ def __init__(
         return Tasks or context IDs) will receive all session events on every
         request. If False, the default behavior of sending only events since the
         last reply from the agent will be used.
+      context_builder: Optional callable that builds ``(message_parts,
+        context_id)`` for the remote A2A request. When set, it replaces
+        ``_construct_message_parts_from_session``. When ``None``, the default
+        session history construction is used (backward compatible).
       config: Optional configuration object.
       use_legacy: If false, send request to the server including the extension
         indicating that the server should use the new implementation.
@@ -236,6 +272,7 @@ def __init__(
     self._a2a_client_factory: Optional[A2AClientFactory] = a2a_client_factory
     self._a2a_request_meta_provider = a2a_request_meta_provider
     self._full_history_when_stateless = full_history_when_stateless
+    self._context_builder = context_builder
     self._config = config or A2aRemoteAgentConfig()
 
     if not use_legacy:
@@ -604,6 +641,18 @@ def _construct_message_parts_from_session(
 
     return message_parts, context_id
 
+  def _build_message_parts_for_request(
+      self, ctx: InvocationContext
+  ) -> tuple[list[A2APart], Optional[str]]:
+    """Build A2A message parts for the outgoing remote request.
+
+    Uses ``context_builder`` when provided; otherwise falls back to
+    ``_construct_message_parts_from_session``.
+    """
+    if self._context_builder is not None:
+      return self._context_builder(ctx, self.name, self._genai_part_converter)
+    return self._construct_message_parts_from_session(ctx)
+
   async def _handle_a2a_response(
       self,
       a2a_response: _compat.A2AClientEvent | A2AMessage,
@@ -827,9 +876,7 @@ async def _run_async_impl(
     # Create A2A request for function response or regular message
     a2a_request = 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 = self._build_message_parts_for_request(ctx)
 
       if not message_parts:
         logger.warning(

tests/unittests/agents/test_remote_a2a_agent.py

@@ -1548,6 +1548,61 @@ def mock_converter(part):
       assert _compat.part_text(parts[0]) == "User question"
       assert _compat.part_text(parts[1]) == "For context:"
 
+  def test_init_accepts_context_builder(self):
+    """RemoteA2aAgent stores an optional context_builder callable."""
+
+    def _builder(ctx, agent_name, genai_part_converter):
+      del ctx, agent_name, genai_part_converter
+      return [], None
+
+    agent = RemoteA2aAgent(
+        name="test_agent",
+        agent_card=create_test_agent_card(),
+        context_builder=_builder,
+    )
+    assert agent._context_builder is _builder
+
+  def test_build_message_parts_uses_custom_context_builder(self):
+    """Custom context_builder replaces default session construction."""
+    custom_part = _compat.make_text_part("only-current-turn")
+    seen = {}
+
+    def _builder(ctx, agent_name, genai_part_converter):
+      seen["ctx"] = ctx
+      seen["agent_name"] = agent_name
+      seen["converter"] = genai_part_converter
+      return [custom_part], "ctx-custom"
+
+    self.agent._context_builder = _builder
+    with patch.object(
+        self.agent, "_construct_message_parts_from_session"
+    ) as mock_default:
+      parts, context_id = self.agent._build_message_parts_for_request(
+          self.mock_context
+      )
+
+    assert parts == [custom_part]
+    assert context_id == "ctx-custom"
+    assert seen["ctx"] is self.mock_context
+    assert seen["agent_name"] == self.agent.name
+    assert seen["converter"] is self.agent._genai_part_converter
+    mock_default.assert_not_called()
+
+  def test_build_message_parts_falls_back_without_context_builder(self):
+    """Without context_builder, default session construction is used."""
+    self.agent._context_builder = None
+    with patch.object(
+        self.agent, "_construct_message_parts_from_session"
+    ) as mock_default:
+      mock_default.return_value = ([], None)
+      parts, context_id = self.agent._build_message_parts_for_request(
+          self.mock_context
+      )
+
+    assert parts == []
+    assert context_id is None
+    mock_default.assert_called_once_with(self.mock_context)
+
   @pytest.mark.asyncio
   async def test_handle_a2a_response_with_task_submitted_and_no_update(self):
     """Test successful A2A response handling with streaming task and no update."""
@@ -2967,6 +3022,72 @@ async def test_run_async_impl_successful_request(self):
                     in mock_event.custom_metadata
                 )
 
+  @pytest.mark.asyncio
+  async def test_run_async_impl_uses_custom_context_builder(self):
+    """_run_async_impl uses context_builder instead of default construction."""
+    custom_part = _compat.make_text_part("from-builder")
+
+    def _builder(ctx, agent_name, genai_part_converter):
+      del ctx, agent_name, genai_part_converter
+      return [custom_part], "builder-context"
+
+    self.agent._context_builder = _builder
+
+    with patch.object(self.agent, "_ensure_resolved"):
+      with patch.object(
+          self.agent, "_create_a2a_request_for_user_function_response"
+      ) as mock_create_func:
+        mock_create_func.return_value = None
+        with patch.object(
+            self.agent, "_construct_message_parts_from_session"
+        ) as mock_default:
+          mock_a2a_client = create_autospec(spec=A2AClient, instance=True)
+          mock_response = _make_stream_message(
+              A2AMessage(
+                  message_id="m1",
+                  role=_compat.ROLE_USER,
+                  parts=[custom_part],
+              )
+          )
+          mock_send_message = AsyncMock()
+          mock_send_message.__aiter__.return_value = [mock_response]
+          mock_a2a_client.send_message.return_value = mock_send_message
+          self.agent._a2a_client = mock_a2a_client
+          self.agent._ensure_resolved.return_value = mock_a2a_client
+
+          mock_event = Event(
+              author=self.agent.name,
+              invocation_id=self.mock_context.invocation_id,
+              branch=self.mock_context.branch,
+          )
+          with patch.object(self.agent, "_handle_a2a_response") as mock_handle:
+            mock_handle.return_value = mock_event
+            with patch(
+                "google.adk.agents.remote_a2a_agent.build_a2a_request_log",
+                return_value="Mock request log",
+            ):
+              with patch(
+                  "google.adk.agents.remote_a2a_agent.build_a2a_response_log",
+                  return_value="Mock response log",
+              ):
+                with patch(
+                    "google.adk.a2a._compat.a2a_to_dict",
+                    return_value={"k": "v"},
+                ):
+                  events = []
+                  async for event in self.agent._run_async_impl(
+                      self.mock_context
+                  ):
+                    events.append(event)
+
+          mock_default.assert_not_called()
+          assert len(events) == 1
+          assert events[0] == mock_event
+          request_meta = mock_event.custom_metadata[
+              A2A_METADATA_PREFIX + "request"
+          ]
+          assert request_meta == {"k": "v"}
+
   @pytest.mark.asyncio
   async def test_run_async_impl_closes_stream_when_abandoned(self):
     """The A2A stream is closed when the caller stops consuming early."""