From 568f7a66677f0c085765b0c4af1c8cd40cf7b722 Mon Sep 17 00:00:00 2001 From: samecho Date: Thu, 29 Oct 2026 18:45:47 -0400 Subject: [PATCH] fix(a2a): recover completed peers when a paused caller resumes Read the recorded peer task before continuing a paused task-mode delegation. Recover its matching finish_task result when the peer completed out of band, without forwarding the peer's private conversation history. Related #7811 --- src/google/adk/a2a/_compat.py | 13 + src/google/adk/a2a/agent/_remote_a2a_agent.py | 106 +++++++++ .../a2a/agent/test_remote_a2a_agent.py | 222 +++++++++++++++++- tests/unittests/a2a/test_compat.py | 17 ++ 4 files changed, 357 insertions(+), 1 deletion(-) diff --git a/src/google/adk/a2a/_compat.py b/src/google/adk/a2a/_compat.py index 681084a87af..eb7db249f7c 100644 --- a/src/google/adk/a2a/_compat.py +++ b/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 # ----------------------------------------------------------------------------- diff --git a/src/google/adk/a2a/agent/_remote_a2a_agent.py b/src/google/adk/a2a/agent/_remote_a2a_agent.py index 1fa999c289c..cb1fb679fd7 100644 --- a/src/google/adk/a2a/agent/_remote_a2a_agent.py +++ b/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 diff --git a/tests/unittests/a2a/agent/test_remote_a2a_agent.py b/tests/unittests/a2a/agent/test_remote_a2a_agent.py index 4d7579d65c2..d2502b4b109 100644 --- a/tests/unittests/a2a/agent/test_remote_a2a_agent.py +++ b/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( diff --git a/tests/unittests/a2a/test_compat.py b/tests/unittests/a2a/test_compat.py index 62538c13268..a3f10f15027 100644 --- a/tests/unittests/a2a/test_compat.py +++ b/tests/unittests/a2a/test_compat.py @@ -28,7 +28,9 @@ import base64 import json +from unittest.mock import create_autospec +from a2a.client import Client from a2a.client.client import ClientConfig from a2a.client.client_factory import ClientFactory from a2a.types import AgentCapabilities @@ -84,6 +86,21 @@ def _task(): ) +async def test_get_task_preserves_the_task_id_and_client_context(): + """The installed SDK receives its native query and the caller's context.""" + client = create_autospec(Client, instance=True) + task = _task() + client.get_task.return_value = task + context = _compat.ClientCallContext(state={'request': 'state'}) + + result = await _compat.get_task(client, task_id='task-1', context=context) + + assert result is task + client.get_task.assert_awaited_once() + assert client.get_task.call_args.args[0].id == 'task-1' + assert client.get_task.call_args.kwargs['context'] is context + + # ----------------------------------------------------------------------------- # make_task # -----------------------------------------------------------------------------