fix(a2a): recover completed peers when a paused caller resumes

src/google/adk/a2a/_compat.py

@@ -60,6 +60,9 @@
 
 from ..utils.context_utils import Aclosing
 
+if TYPE_CHECKING:
+  from a2a.client import Client as A2AClient
+
 # Re-exported: callers construct factories via ``_compat.A2AClientFactory``.
 A2AClientFactory = _ClientFactory
 
@@ -868,6 +871,16 @@ async def send_message(
         yield item
 
 
+async def get_task(
+    client: A2AClient, *, task_id: str, context: Any = None
+) -> Task:
+  """Reads a task using the request type supported by the installed SDK."""
+  request_type = _dynamic_type(
+      "a2a.types", "GetTaskRequest" if IS_A2A_V1 else "TaskQueryParams"
+  )
+  return await client.get_task(request_type(id=task_id), context=context)
+
+
 # -----------------------------------------------------------------------------
 # Client config builder
 # -----------------------------------------------------------------------------

src/google/adk/a2a/agent/_remote_a2a_agent.py

@@ -1974,6 +1974,104 @@ async def _handle_a2a_response_v2(
           branch=ctx.branch,
       )
 
+  async def _read_completed_peer_task(
+      self,
+      ctx: InvocationContext,
+      client: A2AClient,
+      request: A2AMessage,
+      parameters: ParametersConfig,
+  ) -> Optional[Event]:
+    """Reads a paused delegation's result if the peer finished out of band."""
+    if (
+        self.mode != AgentMode.TASK
+        or not request.task_id
+        or request.task_id
+        != self._get_resumable_task_id(ctx, request.context_id)
+    ):
+      return None
+
+    task = await _compat.get_task(
+        client,
+        task_id=request.task_id,
+        context=parameters.client_call_context,
+    )
+    if task.id != request.task_id or task.context_id != request.context_id:
+      raise A2AClientError("Peer returned a different task or context")
+    if task.status.state != _compat.TS_COMPLETED:
+      return None
+
+    if _compat.metadata_get(task.metadata, _NEW_A2A_ADK_INTEGRATION_EXTENSION):
+      event = await self._handle_a2a_response_v2((task, None), ctx)
+    else:
+      event = await self._handle_a2a_response((task, None), ctx)
+    completed_response = (
+        next(
+            (
+                response
+                for response in reversed(event.get_function_responses())
+                if response.name == FINISH_TASK_TOOL_NAME
+                and response.response
+                and response.response.get("result")
+                == FINISH_TASK_SUCCESS_RESULT
+            ),
+            None,
+        )
+        if event is not None
+        else None
+    )
+    if completed_response is None:
+      return None
+    completed_id = completed_response.id
+
+    def output_args(source: Optional[Event]) -> Optional[dict[str, Any]]:
+      if source is None:
+        return None
+      return next(
+          (
+              dict(call.args or {})
+              for call in reversed(source.get_function_calls())
+              if call.name == FINISH_TASK_TOOL_NAME and call.id == completed_id
+          ),
+          None,
+      )
+
+    args = output_args(event)
+    if args is None:
+      # A unary task snapshot can put only the response in its artifacts.
+      for message in reversed(task.history or []):
+        if message.role != _compat.ROLE_AGENT:
+          continue
+        args = output_args(
+            convert_a2a_message_to_event(
+                message, self.name, ctx, self._a2a_part_converter
+            )
+        )
+        if args is not None:
+          break
+    if args is None:
+      raise A2AClientError(
+          "Completed peer task has no matching finish_task call"
+      )
+
+    wrapper_key = get_output_wrapper_key(self.output_schema)
+    output = args[wrapper_key] if wrapper_key and wrapper_key in args else args
+    result = _create_finish_task_event(ctx, self.name, output=output)
+    result.custom_metadata = {
+        A2A_METADATA_PREFIX + "operation": "get_task",
+        A2A_METADATA_PREFIX + "task_id": task.id,
+        A2A_METADATA_PREFIX + "context_id": task.context_id,
+        A2A_METADATA_PREFIX + "request": {"id": task.id},
+        A2A_METADATA_PREFIX
+        + "response": {
+            "id": task.id,
+            "contextId": task.context_id,
+            "status": {
+                "state": _compat.a2a_to_dict(task.status).get("state"),
+            },
+        },
+    }
+    return result
+
   async def _run_async_impl(
       self, ctx: InvocationContext
   ) -> AsyncGenerator[Event, None]:
@@ -2072,6 +2170,14 @@ async def _run_async_impl(
               ctx, a2a_request
           )
 
+        completed_event = await self._read_completed_peer_task(
+            ctx, a2a_client, a2a_request, parameters
+        )
+        if completed_event is not None:
+          should_release_task_control = True
+          yield completed_event
+          return
+
         # TODO: Add support for requested_extension and
         # message_send_configuration once they are supported by the A2A client.
         # A single stateful normalizer per stream so incremental

tests/unittests/a2a/agent/test_remote_a2a_agent.py

@@ -8707,6 +8707,219 @@ async def test_user_input_resumes_paused_remote_task(
   assert not requests[2].task_id
 
 
+def _completed_peer_setup(
+    *, use_v2=False, output_schema=None, call_in_history=False, config=None
+):
+  """Create a paused delegation and its completed peer at the SDK boundary."""
+  from google.adk.sessions import InMemorySessionService
+
+  agent = RemoteA2aAgent(
+      name="currency",
+      mode="task",
+      output_schema=output_schema,
+      config=config,
+      agent_card=create_test_agent_card(name="currency"),
+      context_builder=lambda ctx, name, converter: (
+          [_compat.make_text_part("Continue")],
+          "context",
+      ),
+  )
+  session = Session(
+      id="session",
+      app_name="continuation",
+      user_id="user",
+      events=[
+          _make_dummy_task_trigger_event("delegation", "currency"),
+          Event(
+              author="currency",
+              isolation_scope="delegation",
+              custom_metadata={
+                  "a2a:task_id": "peer-task",
+                  "a2a:context_id": "context",
+                  "a2a:response": {
+                      "status": {"state": "input-required"},
+                  },
+              },
+          ),
+      ],
+  )
+  ctx = InvocationContext(
+      invocation_id="invocation",
+      session=session,
+      session_service=InMemorySessionService(),
+      agent=agent,
+      isolation_scope="delegation",
+  )
+  client = create_autospec(A2AClient, instance=True)
+
+  async def send_message(*args, **kwargs):
+    raise RuntimeError("The peer task is already completed")
+    yield
+
+  client.send_message = Mock(side_effect=send_message)
+  call = convert_genai_part_to_a2a_part(
+      genai_types.Part(
+          function_call=genai_types.FunctionCall(
+              id="finish", name="finish_task", args={"result": "17 CAD"}
+          )
+      )
+  )
+  response = convert_genai_part_to_a2a_part(
+      genai_types.Part(
+          function_response=genai_types.FunctionResponse(
+              id="finish",
+              name="finish_task",
+              response={"result": FINISH_TASK_SUCCESS_RESULT},
+          )
+      )
+  )
+  client.get_task.return_value = _compat.make_task(
+      id="peer-task",
+      context_id="context",
+      metadata=(
+          {remote_a2a_agent._NEW_A2A_ADK_INTEGRATION_EXTENSION: True}
+          if use_v2
+          else None
+      ),
+      status=_compat.make_task_status(
+          _compat.TS_COMPLETED,
+          message=_compat.make_message(
+              message_id="final",
+              role="agent",
+              parts=[response] if call_in_history else [call, response],
+          ),
+      ),
+      history=[
+          _compat.make_message(
+              message_id="finish-call", role="agent", parts=[call]
+          )
+      ]
+      if call_in_history
+      else [],
+  )
+
+  return agent, ctx, client
+
+
+@pytest.mark.parametrize("use_v2", [False, True], ids=["legacy", "v2"])
+@pytest.mark.parametrize("output_schema", [None, str])
+@pytest.mark.parametrize("call_in_history", [False, True])
+async def test_completed_peer_is_read_without_resending_input(
+    use_v2, output_schema, call_in_history
+):
+  """A peer completed out of band supplies the pending delegation's result."""
+  agent, ctx, client = _completed_peer_setup(
+      use_v2=use_v2,
+      output_schema=output_schema,
+      call_in_history=call_in_history,
+  )
+
+  with patch.object(agent, "_ensure_resolved", return_value=client):
+    events = [event async for event in agent.run_async(ctx)]
+
+  expected = "17 CAD" if output_schema is str else {"result": "17 CAD"}
+  assert [event.output for event in events if event.output is not None] == [
+      expected
+  ]
+  client.send_message.assert_not_called()
+  client.get_task.assert_awaited_once()
+  assert client.get_task.call_args.args[0].id == "peer-task"
+
+
+@pytest.mark.parametrize("field", ["id", "context_id"])
+async def test_completed_peer_from_another_task_is_not_used(field):
+  """A lookup for one delegation cannot supply another task's result."""
+  agent, ctx, client = _completed_peer_setup()
+  setattr(client.get_task.return_value, field, "different")
+
+  with patch.object(agent, "_ensure_resolved", return_value=client):
+    events = [event async for event in agent.run_async(ctx)]
+
+  assert not any(event.output is not None for event in events)
+  assert any(
+      "different task or context" in (event.error_message or "")
+      for event in events
+  )
+  client.send_message.assert_not_called()
+
+
+async def test_completed_peer_does_not_take_output_from_user_history():
+  """Only the peer's own finish_task call can supply the completed result."""
+  agent, ctx, client = _completed_peer_setup(call_in_history=True)
+  client.get_task.return_value.history[0].role = _compat.ROLE_USER
+
+  with patch.object(agent, "_ensure_resolved", return_value=client):
+    events = [event async for event in agent.run_async(ctx)]
+
+  assert not any(event.output is not None for event in events)
+  assert any(
+      "no matching finish_task call" in (event.error_message or "")
+      for event in events
+  )
+  client.send_message.assert_not_called()
+
+
+async def test_peer_lookup_failure_does_not_resend_input():
+  """An unsuccessful task lookup must not repeat a confirmation message."""
+  agent, ctx, client = _completed_peer_setup()
+  client.get_task.side_effect = RuntimeError("Task lookup unavailable")
+
+  with patch.object(agent, "_ensure_resolved", return_value=client):
+    events = [event async for event in agent.run_async(ctx)]
+
+  assert any(
+      "Task lookup unavailable" in (event.error_message or "")
+      for event in events
+  )
+  client.send_message.assert_not_called()
+
+
+async def test_peer_lookup_does_not_persist_private_task_history():
+  """Out-of-band user input stays out of the caller's persisted events."""
+  agent, ctx, client = _completed_peer_setup(call_in_history=True)
+  client.get_task.return_value.history.append(
+      _compat.make_message(
+          message_id="private",
+          role="user",
+          parts=[_compat.make_data_part(data={"token": "private-test-token"})],
+      )
+  )
+
+  with patch.object(agent, "_ensure_resolved", return_value=client):
+    events = [event async for event in agent.run_async(ctx)]
+
+  assert [event.output for event in events if event.output is not None] == [
+      {"result": "17 CAD"}
+  ]
+  assert all(
+      "private-test-token" not in event.model_dump_json() for event in events
+  )
+
+
+async def test_peer_lookup_keeps_the_before_request_call_context():
+  """The lookup carries the same caller context as a normal request."""
+  call_context = _compat.ClientCallContext(state={"request": "state"})
+
+  async def before_request(ctx, request, parameters):
+    return request, ParametersConfig(client_call_context=call_context)
+
+  agent, ctx, client = _completed_peer_setup(
+      config=A2aRemoteAgentConfig(
+          request_interceptors=[
+              RequestInterceptor(before_request=before_request)
+          ]
+      )
+  )
+
+  with patch.object(agent, "_ensure_resolved", return_value=client):
+    events = [event async for event in agent.run_async(ctx)]
+
+  assert [event.output for event in events if event.output is not None] == [
+      {"result": "17 CAD"}
+  ]
+  assert client.get_task.call_args.kwargs["context"] is call_context
+
+
 @pytest.mark.parametrize(
     "initial_state",
     [
@@ -8805,9 +9018,15 @@ async def send_message(client, *, request, **kwargs):
         )
     )
 
+  client = create_autospec(A2AClient, instance=True)
+  client.get_task.return_value = _compat.make_task(
+      id="current-task",
+      context_id="context",
+      status=_compat.make_task_status(_compat.TS_INPUT_REQUIRED),
+  )
   try:
     with (
-        patch.object(agent, "_ensure_resolved", return_value=object()),
+        patch.object(agent, "_ensure_resolved", return_value=client),
         patch.object(_compat, "send_message", side_effect=send_message),
     ):
       async for _ in agent.run_async(ctx):
@@ -8817,6 +9036,7 @@ async def send_message(client, *, request, **kwargs):
 
   assert len(requests) == 1
   assert (requests[0].task_id or None) == expected_task_id
+  assert client.get_task.await_count == (1 if expected_task_id else 0)
 
 
 def _bare_task_stream(