From 0cf565d049a86b60901b081192e0bca7f73b28b6 Mon Sep 17 00:00:00 2001 From: maximilliangrand <214999687+maximilliangrand@users.noreply.github.com> Date: Sun, 27 Sep 2026 14:07:06 +0200 Subject: [PATCH] fix: resume paused A2A tasks on ordinary user replies --- src/google/adk/a2a/_compat.py | 2 + src/google/adk/a2a/agent/_remote_a2a_agent.py | 104 ++- src/google/adk/a2a/converters/to_adk_event.py | 36 +- .../a2a/agent/test_remote_a2a_agent.py | 687 +++++++++++++++++- tests/unittests/a2a/converters/test_to_adk.py | 99 +++ 5 files changed, 922 insertions(+), 6 deletions(-) diff --git a/src/google/adk/a2a/_compat.py b/src/google/adk/a2a/_compat.py index b2bba8e1d36..bd04c59000a 100644 --- a/src/google/adk/a2a/_compat.py +++ b/src/google/adk/a2a/_compat.py @@ -148,6 +148,7 @@ def _proto_to_dict(msg: Any) -> dict[str, Any]: TS_INPUT_REQUIRED = TaskState.Value("TASK_STATE_INPUT_REQUIRED") TS_AUTH_REQUIRED = TaskState.Value("TASK_STATE_AUTH_REQUIRED") TS_CANCELED = TaskState.Value("TASK_STATE_CANCELED") + TS_REJECTED = TaskState.Value("TASK_STATE_REJECTED") TP_JSONRPC = TransportProtocol.JSONRPC TP_HTTP_JSON = TransportProtocol.HTTP_JSON @@ -164,6 +165,7 @@ def _proto_to_dict(msg: Any) -> dict[str, Any]: TS_INPUT_REQUIRED = TaskState.input_required TS_AUTH_REQUIRED = TaskState.auth_required TS_CANCELED = TaskState.canceled + TS_REJECTED = TaskState.rejected TP_JSONRPC = getattr(TransportProtocol, "jsonrpc") TP_HTTP_JSON = getattr(TransportProtocol, "http_json") diff --git a/src/google/adk/a2a/agent/_remote_a2a_agent.py b/src/google/adk/a2a/agent/_remote_a2a_agent.py index 397402df855..a8201ae04c0 100644 --- a/src/google/adk/a2a/agent/_remote_a2a_agent.py +++ b/src/google/adk/a2a/agent/_remote_a2a_agent.py @@ -33,6 +33,7 @@ from a2a.types import AgentCard from a2a.types import Message as A2AMessage from a2a.types import Part as A2APart +from a2a.types import Task as A2ATask from a2a.types import TaskArtifactUpdateEvent as A2ATaskArtifactUpdateEvent from a2a.types import TaskState from a2a.types import TaskStatusUpdateEvent as A2ATaskStatusUpdateEvent @@ -77,6 +78,7 @@ from ..converters.part_converter import convert_genai_part_to_a2a_part from ..converters.part_converter import GenAIPartToA2APartConverter from ..converters.to_adk_event import _create_mock_function_call_for_required_user_input +from ..converters.to_adk_event import _TASK_RESPONSE_BOUNDARY_STATES from ..converters.to_adk_event import MOCK_FUNCTION_CALL_FOR_REQUIRED_USER_AUTH from ..converters.to_adk_event import MOCK_FUNCTION_CALL_FOR_REQUIRED_USER_INPUT from ..experimental import a2a_experimental @@ -147,6 +149,18 @@ _compat.TS_WORKING, }) + +def _task_response_metadata(task: A2ATask) -> dict[str, Any]: + """Records the latest task snapshot, including contentless status updates.""" + metadata = { + A2A_METADATA_PREFIX + "task_id": task.id, + A2A_METADATA_PREFIX + "response": _compat.a2a_to_dict(task), + } + if task.context_id: + metadata[A2A_METADATA_PREFIX + "context_id"] = task.context_id + return metadata + + _RESULT_KEY = "result" # Top-level keys of a serialized AuthConfig, the shape an adk_request_credential @@ -1381,6 +1395,73 @@ def _build_message_parts_for_request( return self._context_builder(ctx, self.name, self._genai_part_converter) return self._construct_message_parts_from_session(ctx) + def _get_resumable_task_id( + self, ctx: InvocationContext, context_id: Optional[str] + ) -> Optional[str]: + """Returns the latest paused task in the outgoing message's context.""" + if not context_id or not self._has_default_response_persistence(): + return None + task_scope = ctx.isolation_scope if self.mode == "task" else None + for event in reversed(ctx.session.events): + if task_scope and event.isolation_scope != task_scope: + # Do not resume a task from an earlier delegation to the same agent. + if any(fc.id == task_scope for fc in event.get_function_calls()): + return None + continue + metadata = event.custom_metadata or {} + if ( + event.author == self.name + and event.error_message + and metadata.get(A2A_METADATA_PREFIX + "request") + ): + # A failed request may already have consumed the recorded pause. + # Its old task state is no longer sufficient to infer a continuation. + return None + if not self._is_remote_response(event): + continue + if metadata.get(A2A_METADATA_PREFIX + "context_id") != context_id: + return None + response = metadata.get(A2A_METADATA_PREFIX + "response") + status = response.get("status") if isinstance(response, dict) else None + # Session metadata contains the serialized wire state. A2A 0.3 uses + # hyphenated names, while A2A 1.x uses protobuf enum names. + if not isinstance(status, dict) or status.get("state") not in ( + "input-required", + "auth-required", + "TASK_STATE_INPUT_REQUIRED", + "TASK_STATE_AUTH_REQUIRED", + ): + return None + task_id = metadata.get(A2A_METADATA_PREFIX + "task_id") + return task_id if isinstance(task_id, str) else None + return None + + def _has_default_response_persistence(self) -> bool: + """Whether recorded responses can safely determine task continuation. + + Inbound converters and after-request hooks may suppress a terminal response, + leaving an older pause in history. With these overrides, callers retain + control of task routing through function responses or request interceptors. + """ + if self._a2a_part_converter is not convert_a2a_part_to_genai_part: + return False + for name in ( + "a2a_task_converter", + "a2a_status_update_converter", + "a2a_message_converter", + "a2a_artifact_update_converter", + "a2a_part_converter", + ): + if ( + getattr(self._config, name) + is not A2aRemoteAgentConfig.model_fields[name].default + ): + return False + return not any( + interceptor.after_request is not None + for interceptor in self._config.request_interceptors or [] + ) + async def _handle_a2a_response( self, a2a_response: _compat.A2AClientEvent | A2AMessage, @@ -1447,6 +1528,17 @@ async def _handle_a2a_response( for part in event.content.parts or []: part.thought = True _add_mock_function_call(event, update.status.state) + elif ( + isinstance(update, A2ATaskStatusUpdateEvent) + and update.status.state in _TASK_RESPONSE_BOUNDARY_STATES + ): + # A bare pause or terminal status is still the end of an interaction. + # Persist it so the next user turn sees the current task state. + event = Event( + author=self.name, + invocation_id=ctx.invocation_id, + branch=ctx.branch, + ) elif isinstance(update, A2ATaskArtifactUpdateEvent): # This is a streaming task artifact update. # Convert only the parts carried by this update. Converting the @@ -1667,10 +1759,12 @@ async def _run_async_impl( parts=message_parts, role=_compat.ROLE_USER, context_id=context_id or session_id, + task_id=self._get_resumable_task_id(ctx, context_id), ) logger.debug(build_a2a_request_log(a2a_request)) + last_task = None try: intercepted_request, parameters = ( await execute_before_request_interceptors( @@ -1697,7 +1791,6 @@ async def _run_async_impl( # status/artifact updates are aggregated into a running task (matching the # 0.3.x client behavior). normalize_stream_item = _compat.make_stream_normalizer() - last_task = None async with Aclosing( _compat.send_message( a2a_client, @@ -1832,6 +1925,11 @@ async def _run_async_impl( error_message=task_error_message, invocation_id=ctx.invocation_id, branch=ctx.branch, + custom_metadata={ + **_task_response_metadata(last_task), + A2A_METADATA_PREFIX + + "request": _compat.a2a_to_dict(a2a_request), + }, ) except _compat.A2A_HTTP_ERRORS as e: @@ -1844,6 +1942,8 @@ async def _run_async_impl( A2A_METADATA_PREFIX + "error": error_message, A2A_METADATA_PREFIX + "status_code": str(status_code), } + if last_task is not None: + custom_metadata.update(_task_response_metadata(last_task)) if a2a_request: custom_metadata[A2A_METADATA_PREFIX + "request"] = ( _compat.a2a_to_dict(a2a_request) @@ -1864,6 +1964,8 @@ async def _run_async_impl( custom_metadata = { A2A_METADATA_PREFIX + "error": error_message, } + if last_task is not None: + custom_metadata.update(_task_response_metadata(last_task)) if a2a_request: custom_metadata[A2A_METADATA_PREFIX + "request"] = ( _compat.a2a_to_dict(a2a_request) diff --git a/src/google/adk/a2a/converters/to_adk_event.py b/src/google/adk/a2a/converters/to_adk_event.py index 1d713be7847..d01773d0027 100644 --- a/src/google/adk/a2a/converters/to_adk_event.py +++ b/src/google/adk/a2a/converters/to_adk_event.py @@ -55,6 +55,17 @@ "mock_function_call_for_required_user_auth" ) +# These states finish the current interaction, including a deliberate pause. +# Unlike progress updates, a bare boundary still needs a persisted response. +_TASK_RESPONSE_BOUNDARY_STATES = frozenset({ + _compat.TS_INPUT_REQUIRED, + _compat.TS_AUTH_REQUIRED, + _compat.TS_COMPLETED, + _compat.TS_FAILED, + _compat.TS_CANCELED, + _compat.TS_REJECTED, +}) + A2AMessageToEventConverter = Callable[ [ Message, @@ -196,11 +207,21 @@ def _create_event( usage_metadata: Any = None, error_code: Any = None, citation_metadata: Any = None, + allow_empty_event: bool = False, ) -> Optional[Event]: """Creates an ADK event from parts and metadata.""" event_actions = actions or EventActions() - if not output_parts and not event_actions.model_dump( - exclude_none=True, exclude_defaults=True + if ( + not allow_empty_event + and not output_parts + and event_actions == EventActions() + and not any(( + grounding_metadata, + custom_metadata, + usage_metadata, + error_code, + citation_metadata, + )) ): return None @@ -496,11 +517,13 @@ def convert_a2a_task_to_event( output_parts: list[genai_types.Part] = [] long_running_function_ids: set[str] = set() metadata_fields: dict[str, Any] = {} + has_source_parts = False status_message = _compat.normalize_message(a2a_task.status.message) if a2a_task.artifacts: artifact_parts = [ part for artifact in a2a_task.artifacts for part in artifact.parts ] + has_source_parts = bool(artifact_parts) for artifact in a2a_task.artifacts: event_actions = _merge_event_actions( event_actions, _extract_event_actions(artifact.metadata) @@ -523,6 +546,7 @@ def convert_a2a_task_to_event( parts, ids = _convert_a2a_parts_to_adk_parts( status_message.parts, part_converter ) + has_source_parts = has_source_parts or bool(status_message.parts) output_parts.extend(parts) long_running_function_ids.update(ids) elif status_message and not metadata_fields: @@ -550,6 +574,10 @@ def convert_a2a_task_to_event( author, event_actions, long_running_function_ids, + allow_empty_event=( + not has_source_parts + and a2a_task.status.state in _TASK_RESPONSE_BOUNDARY_STATES + ), **metadata_fields, ) @@ -655,6 +683,10 @@ def convert_a2a_status_update_to_event( author, event_actions, long_running_function_ids, + allow_empty_event=( + (status_message is None or not status_message.parts) + and a2a_status_update.status.state in _TASK_RESPONSE_BOUNDARY_STATES + ), **metadata_fields, ) except Exception as e: diff --git a/tests/unittests/a2a/agent/test_remote_a2a_agent.py b/tests/unittests/a2a/agent/test_remote_a2a_agent.py index 5a7d5358b5b..756c10187b6 100644 --- a/tests/unittests/a2a/agent/test_remote_a2a_agent.py +++ b/tests/unittests/a2a/agent/test_remote_a2a_agent.py @@ -29,7 +29,6 @@ from a2a.types import AgentCard from a2a.types import AgentInterface from a2a.types import AgentSkill -from a2a.types import Artifact from a2a.types import Message as A2AMessage from a2a.types import Task as A2ATask from a2a.types import TaskArtifactUpdateEvent @@ -2282,6 +2281,7 @@ async def test_handle_a2a_response_with_task_status_update_no_message(self): """Test handling of a task status update with no message.""" mock_a2a_task = Mock(spec=A2ATask) mock_a2a_task.id = "task-123" + mock_a2a_task.context_id = "context-123" mock_update = Mock(spec=TaskStatusUpdateEvent) mock_update.status = Mock(A2ATaskStatus) @@ -2292,7 +2292,9 @@ async def test_handle_a2a_response_with_task_status_update_no_message(self): (mock_a2a_task, mock_update), self.mock_context ) - assert result is None + assert result.content is None + assert result.is_final_response() + assert result.custom_metadata["a2a:task_id"] == "task-123" @pytest.mark.asyncio async def test_handle_a2a_response_filters_thought_parts_from_completed_task( @@ -3422,6 +3424,7 @@ async def test_handle_a2a_response_with_task_status_update_no_message(self): """Test handling of a task status update with no message.""" mock_a2a_task = Mock(spec=A2ATask) mock_a2a_task.id = "task-123" + mock_a2a_task.context_id = "context-123" mock_update = Mock(spec=TaskStatusUpdateEvent) mock_update.status = Mock(A2ATaskStatus) @@ -3432,7 +3435,9 @@ async def test_handle_a2a_response_with_task_status_update_no_message(self): (mock_a2a_task, mock_update), self.mock_context ) - assert result is None + assert result.content is None + assert result.is_final_response() + assert result.custom_metadata["a2a:task_id"] == "task-123" @pytest.mark.asyncio async def test_handle_a2a_response_with_artifact_update(self): @@ -7589,3 +7594,679 @@ async def test_run_async_impl_keeps_task_control_on_auth_request(self): assert events == [auth_request_event] ctx.set_agent_state.assert_not_called() + + +async def _send_remote_task_followups( + initial_state, followup_parts, *, bare_completion=False, repeat_pause=False +): + """Run three public Runner turns against a scripted A2A transport.""" + from google.adk.runners import Runner + from google.adk.sessions import InMemorySessionService + + agent = RemoteA2aAgent( + name="currency", agent_card=create_test_agent_card(name="currency") + ) + requests = [] + + async def send_message(client, *, request, **kwargs): + del client, kwargs + requests.append(request) + state = ( + initial_state + if len(requests) == 1 or (repeat_pause and len(requests) == 2) + else _compat.TS_COMPLETED + ) + task = _compat.make_task( + id="currency-task", + context_id="currency-context", + status=_compat.make_task_status( + state, + message=_compat.make_message( + message_id="reply", + role="agent", + parts=[_compat.make_text_part("Currency response")], + ), + ), + ) + if bare_completion and len(requests) == 2: + update = _compat.make_task_status_update_event( + task_id=task.id, + context_id=task.context_id, + status=_compat.make_task_status(_compat.TS_COMPLETED), + final=True, + ) + if _compat.IS_A2A_V1: + from a2a.types import StreamResponse + + raw = StreamResponse() + raw.status_update.CopyFrom(update) + yield raw + else: + yield task, update + else: + yield _make_stream_task(task) + + runner = Runner( + app_name="continuation", + agent=agent, + session_service=InMemorySessionService(), + auto_create_session=True, + ) + try: + with ( + patch.object(agent, "_ensure_resolved", return_value=object()), + patch.object(_compat, "send_message", side_effect=send_message), + ): + for parts in ( + [genai_types.Part(text="Convert ten USD")], + followup_parts, + [genai_types.Part(text="Start another conversion")], + ): + async for _ in runner.run_async( + user_id="user", + session_id="session", + new_message=genai_types.Content(role="user", parts=parts), + ): + pass + finally: + await runner.close() + return requests + + +@pytest.mark.parametrize( + "initial_state", [_compat.TS_INPUT_REQUIRED, _compat.TS_AUTH_REQUIRED] +) +@pytest.mark.parametrize( + "followup_parts", + [ + [genai_types.Part(text="CAD")], + [ + genai_types.Part.from_bytes( + data=b"CAD", mime_type="application/octet-stream" + ) + ], + ], + ids=["text", "attachment"], +) +async def test_user_input_resumes_paused_remote_task( + initial_state, followup_parts +): + """Text and attachments continue a paused task, not a new task in its context.""" + requests = await _send_remote_task_followups(initial_state, followup_parts) + + assert not requests[0].task_id + assert requests[1].context_id == "currency-context" + assert requests[1].task_id == "currency-task" + assert requests[1].parts + # Completion ends that task even though the remote context is reused. + assert requests[2].context_id == "currency-context" + assert not requests[2].task_id + + +@pytest.mark.parametrize( + "initial_state", + [ + _compat.TS_COMPLETED, + _compat.TS_FAILED, + _compat.TS_CANCELED, + _compat.TS_SUBMITTED, + _compat.TS_WORKING, + _compat.TS_UNKNOWN, + ], +) +async def test_new_message_does_not_resume_remote_task_without_a_pause( + initial_state, +): + """An existing context alone does not make a user message a task continuation.""" + requests = await _send_remote_task_followups( + initial_state, [genai_types.Part(text="Start another conversion")] + ) + + assert requests[1].context_id == "currency-context" + assert not requests[1].task_id + + +@pytest.mark.parametrize( + "scope, latest_state, response_context, expected_task_id", + [ + ("current", "input-required", "context", "current-task"), + ("current", "auth-required", "context", "current-task"), + ("current", "completed", "context", None), + ("current", "input-required", "different-context", None), + ("new-delegation", "input-required", "context", None), + ], +) +async def test_remote_task_continuation_respects_latest_response_and_scope( + scope, latest_state, response_context, expected_task_id +): + """A continuation cannot select a stale task, sibling task, or another context.""" + from google.adk.sessions import InMemorySessionService + + agent = RemoteA2aAgent( + name="currency", + mode="task", + agent_card=create_test_agent_card(name="currency"), + context_builder=lambda ctx, name, converter: ( + [_compat.make_text_part("Continue")], + "context", + ), + ) + + def response(task_id, state, task_scope, context_id="context"): + return Event( + author="currency", + isolation_scope=task_scope, + custom_metadata={ + "a2a:task_id": task_id, + "a2a:context_id": context_id, + "a2a:response": {"status": {"state": state}}, + }, + ) + + session = Session( + id="session", + app_name="continuation", + user_id="user", + events=[ + response("stale-task", "input-required", "current"), + _make_dummy_task_trigger_event("new-delegation", "currency"), + response("current-task", latest_state, "current", response_context), + response("sibling-task", "input-required", "sibling"), + Event( + author="user", + invocation_id="invocation", + content=genai_types.Content( + role="user", parts=[genai_types.Part(text="Continue")] + ), + ), + ], + ) + ctx = InvocationContext( + invocation_id="invocation", + session=session, + session_service=InMemorySessionService(), + agent=agent, + isolation_scope=scope, + ) + requests = [] + + async def send_message(client, *, request, **kwargs): + del client, kwargs + requests.append(request) + yield _make_stream_message( + _compat.make_message( + message_id="reply", + role="agent", + parts=[_compat.make_text_part("Response")], + ) + ) + + try: + with ( + patch.object(agent, "_ensure_resolved", return_value=object()), + patch.object(_compat, "send_message", side_effect=send_message), + ): + async for _ in agent.run_async(ctx): + pass + finally: + await agent.cleanup() + + assert len(requests) == 1 + assert (requests[0].task_id or None) == expected_task_id + + +def _bare_task_stream(state, *, use_v2=False, artifact_text=None): + """Build SDK-native status updates with their accumulated task snapshots.""" + metadata = ( + {remote_a2a_agent._NEW_A2A_ADK_INTEGRATION_EXTENSION: True} + if use_v2 + else None + ) + task = _compat.make_task( + id="currency-task", + context_id="currency-context", + status=_compat.make_task_status(_compat.TS_WORKING), + metadata=metadata, + ) + responses = [_make_stream_task(copy.deepcopy(task))] + if artifact_text is not None: + artifact = _compat.make_artifact( + artifact_id="answer", + parts=[_compat.make_text_part(artifact_text)], + ) + update = TaskArtifactUpdateEvent( + task_id=task.id, + context_id=task.context_id, + artifact=artifact, + last_chunk=True, + ) + if _compat.IS_A2A_V1: + from a2a.types import StreamResponse + + responses.append(StreamResponse(artifact_update=update)) + else: + task.artifacts = [artifact] + responses.append((task.model_copy(deep=True), update)) + status = _compat.make_task_status(state) + update = _compat.make_task_status_update_event( + task_id=task.id, + context_id=task.context_id, + status=status, + final=state not in remote_a2a_agent._UNFINISHED_TASK_STATES, + ) + if _compat.IS_A2A_V1: + from a2a.types import StreamResponse + + responses.append(StreamResponse(status_update=update)) + else: + task.status = status + responses.append((task, update)) + return responses + + +async def _run_remote_task_responses( + responses, *, same_invocation=False, plugins=None, config=None +): + """Exercise the public Runner with a serialized session on each user turn.""" + from google.adk.runners import Runner + from google.adk.sessions import InMemorySessionService + + class SerializedSessionService(InMemorySessionService): + + async def get_session(self, **kwargs): + session = await super().get_session(**kwargs) + return ( + Session.model_validate_json(session.model_dump_json()) + if session is not None + else None + ) + + agent = RemoteA2aAgent( + name="currency", + agent_card=create_test_agent_card(name="currency"), + config=config, + context_builder=( + lambda ctx, name, converter: ( + [_compat.make_text_part(ctx.user_content.parts[0].text)], + "currency-context", + ) + ) + if same_invocation + else None, + ) + requests = [] + turns = [] + service = SerializedSessionService() + + async def send_message(client, *, request, **kwargs): + del client, kwargs + script = responses[len(requests)] + requests.append(request) + for response in script: + if isinstance(response, Exception): + raise response + yield response + + runner = Runner( + app_name="continuation", + agent=agent, + session_service=service, + auto_create_session=True, + plugins=plugins, + ) + try: + with ( + patch.object(agent, "_ensure_resolved", return_value=object()), + patch.object(_compat, "send_message", side_effect=send_message), + ): + for text in ("Convert ten USD", "CAD", "Start another conversion"): + turns.append([ + event + async for event in runner.run_async( + user_id="user", + session_id="session", + invocation_id="same-invocation" if same_invocation else None, + new_message=genai_types.Content( + role="user", parts=[genai_types.Part(text=text)] + ), + ) + ]) + session = await service.get_session( + app_name="continuation", user_id="user", session_id="session" + ) + finally: + await runner.close() + return requests, turns, session + + +@pytest.mark.parametrize("use_v2", [False, True], ids=["legacy", "v2"]) +@pytest.mark.parametrize("same_invocation", [False, True]) +@pytest.mark.parametrize( + "pause", [_compat.TS_INPUT_REQUIRED, _compat.TS_AUTH_REQUIRED] +) +async def test_bare_task_boundaries_survive_session_reload( + use_v2, same_invocation, pause +): + requests, turns, session = await _run_remote_task_responses( + [ + _bare_task_stream(pause, use_v2=use_v2), + _bare_task_stream(_compat.TS_COMPLETED, use_v2=use_v2), + _bare_task_stream(_compat.TS_COMPLETED, use_v2=use_v2), + ], + same_invocation=same_invocation, + ) + assert [r.task_id or None for r in requests] == [None, "currency-task", None] + assert all(turn[-1].content is None for turn in turns) + assert all(turn[-1].is_final_response() for turn in turns) + assert session.events[-1].custom_metadata["a2a:response"]["status"] == ( + turns[-1][-1].custom_metadata["a2a:response"]["status"] + ) + + +@pytest.mark.parametrize("use_v2", [False, True]) +async def test_repeated_bare_pauses_continue_the_same_task(use_v2): + requests, _, _ = await _run_remote_task_responses([ + _bare_task_stream(_compat.TS_INPUT_REQUIRED, use_v2=use_v2), + _bare_task_stream(_compat.TS_AUTH_REQUIRED, use_v2=use_v2), + _bare_task_stream(_compat.TS_COMPLETED, use_v2=use_v2), + ]) + assert [r.task_id or None for r in requests] == [ + None, + "currency-task", + "currency-task", + ] + + +async def test_plugin_handled_turn_does_not_consume_remote_pause(): + from google.adk.plugins.base_plugin import BasePlugin + + class SkipSecondTurn(BasePlugin): + + async def before_run_callback(self, *, invocation_context): + if invocation_context.user_content.parts[0].text == "CAD": + return genai_types.Content( + role="model", parts=[genai_types.Part(text="Please try again")] + ) + + requests, turns, _ = await _run_remote_task_responses( + [ + _bare_task_stream(_compat.TS_INPUT_REQUIRED), + _bare_task_stream(_compat.TS_COMPLETED), + ], + plugins=[SkipSecondTurn(name="skip_second")], + ) + assert turns[1][0].content.parts[0].text == "Please try again" + assert [r.task_id or None for r in requests] == [None, "currency-task"] + + +@pytest.mark.parametrize( + "second_response", + [ + [ + _make_stream_task( + _compat.make_task( + id="currency-task", + context_id="currency-context", + status=_compat.make_task_status(_compat.TS_WORKING), + ) + ) + ], + [httpx.ReadError("Connection lost")], + [RuntimeError("Invalid remote response")], + ], + ids=["unfinished", "http-error", "response-error"], +) +async def test_failed_continuation_does_not_reuse_an_old_pause(second_response): + requests, turns, _ = await _run_remote_task_responses([ + _bare_task_stream(_compat.TS_INPUT_REQUIRED), + second_response, + _bare_task_stream(_compat.TS_COMPLETED), + ]) + assert [r.task_id or None for r in requests] == [None, "currency-task", None] + assert turns[1][-1].error_message + + +@pytest.mark.parametrize("use_v2", [False, True]) +async def test_bare_completion_keeps_preceding_artifact_output(use_v2): + requests, turns, session = await _run_remote_task_responses([ + _bare_task_stream(_compat.TS_INPUT_REQUIRED, use_v2=use_v2), + _bare_task_stream( + _compat.TS_COMPLETED, use_v2=use_v2, artifact_text="Ten CAD" + ), + _bare_task_stream(_compat.TS_COMPLETED, use_v2=use_v2), + ]) + assert not requests[2].task_id + assert [ + part.text + for event in turns[1] + if event.content + for part in event.content.parts + if part.text + ] == ["Ten CAD"] + assert turns[1][-1].content is None + assert any( + event.content and any(p.text == "Ten CAD" for p in event.content.parts) + for event in session.events + ) + + +async def test_filtered_completion_does_not_infer_continuation_from_old_pause(): + async def filter_completion(ctx, response, event): + del ctx + if ( + isinstance(response, tuple) + and response[0].status.state == _compat.TS_COMPLETED + ): + return None + return event + + config = A2aRemoteAgentConfig( + request_interceptors=[RequestInterceptor(after_request=filter_completion)] + ) + requests, turns, _ = await _run_remote_task_responses( + [ + _bare_task_stream(_compat.TS_INPUT_REQUIRED), + _bare_task_stream(_compat.TS_COMPLETED), + _bare_task_stream(_compat.TS_COMPLETED), + ], + config=config, + ) + assert turns[0][-1].custom_metadata["a2a:task_id"] == "currency-task" + assert all(not request.task_id for request in requests) + assert not any( + event.custom_metadata.get("a2a:response", {}) + .get("status", {}) + .get("state") + in ("completed", "TASK_STATE_COMPLETED") + for turn in turns[1:] + for event in turn + if event.custom_metadata + ) + + +@pytest.mark.parametrize( + "converter_name", + [ + "a2a_task_converter", + "a2a_status_update_converter", + "a2a_message_converter", + "a2a_artifact_update_converter", + "a2a_part_converter", + ], +) +async def test_custom_inbound_converters_keep_explicit_task_routing( + converter_name, +): + default = A2aRemoteAgentConfig.model_fields[converter_name].default + + def converter(*args): + return default(*args) + + requests, _, _ = await _run_remote_task_responses( + [ + _bare_task_stream(_compat.TS_INPUT_REQUIRED, use_v2=True), + _bare_task_stream(_compat.TS_COMPLETED, use_v2=True), + _bare_task_stream(_compat.TS_COMPLETED, use_v2=True), + ], + config=A2aRemoteAgentConfig(**{converter_name: converter}), + ) + assert all(not request.task_id for request in requests) + + +async def test_before_request_hook_keeps_automatic_task_continuation(): + async def before_request(ctx, request, parameters): + del ctx + return request, parameters + + requests, _, _ = await _run_remote_task_responses( + [ + _bare_task_stream(_compat.TS_INPUT_REQUIRED), + _bare_task_stream(_compat.TS_COMPLETED), + _bare_task_stream(_compat.TS_COMPLETED), + ], + config=A2aRemoteAgentConfig( + request_interceptors=[ + RequestInterceptor(before_request=before_request) + ] + ), + ) + assert [r.task_id or None for r in requests] == [None, "currency-task", None] + + +@pytest.mark.parametrize("use_v2", [False, True]) +async def test_agent_tool_keeps_artifact_before_bare_completion(use_v2): + from google.adk.sessions import InMemorySessionService + from google.adk.tools.agent_tool import AgentTool + from google.adk.tools.tool_context import ToolContext + + agent = RemoteA2aAgent( + name="currency", agent_card=create_test_agent_card(name="currency") + ) + session_service = InMemorySessionService() + session = await session_service.create_session( + app_name="parent", user_id="user" + ) + ctx = InvocationContext( + invocation_id="invocation", + session=session, + session_service=session_service, + agent=agent, + ) + + async def send_message(client, **kwargs): + del client, kwargs + for response in _bare_task_stream( + _compat.TS_COMPLETED, use_v2=use_v2, artifact_text="Ten CAD" + ): + yield response + + with ( + patch.object(agent, "_ensure_resolved", return_value=object()), + patch.object(_compat, "send_message", side_effect=send_message), + ): + result = await AgentTool(agent=agent).run_async( + args={"request": "Convert ten USD to CAD"}, + tool_context=ToolContext(ctx), + ) + assert result == "Ten CAD" + + +@pytest.mark.parametrize( + "override", ["legacy-part", "v2-status", "after-request"] +) +def test_response_filters_preserve_explicit_function_response_routing(override): + from google.adk.sessions import InMemorySessionService + + kwargs = {} + if override == "legacy-part": + kwargs["a2a_part_converter"] = lambda part: None + elif override == "v2-status": + kwargs["config"] = A2aRemoteAgentConfig( + a2a_status_update_converter=lambda *args: None + ) + else: + + async def filter_response(ctx, response, event): + return None + + kwargs["config"] = A2aRemoteAgentConfig( + request_interceptors=[RequestInterceptor(after_request=filter_response)] + ) + agent = RemoteA2aAgent( + name="currency", + agent_card=create_test_agent_card(name="currency"), + **kwargs, + ) + session = Session( + id="session", + app_name="continuation", + user_id="user", + events=[ + Event( + author="currency", + content=genai_types.Content( + role="model", + parts=[ + genai_types.Part( + function_call=genai_types.FunctionCall( + id="input", name="select_currency", args={} + ) + ) + ], + ), + custom_metadata={ + "a2a:task_id": "currency-task", + "a2a:context_id": "currency-context", + "a2a:response": {"status": {"state": "input-required"}}, + }, + ), + Event( + author="user", + content=genai_types.Content( + role="user", + parts=[ + genai_types.Part( + function_response=genai_types.FunctionResponse( + id="input", + name="select_currency", + response={"value": "CAD"}, + ) + ) + ], + ), + ), + ], + ) + ctx = InvocationContext( + invocation_id="invocation", + session=session, + session_service=InMemorySessionService(), + agent=agent, + ) + assert agent._get_resumable_task_id(ctx, "currency-context") is None + request = agent._create_a2a_request_for_user_function_response(ctx) + assert request.task_id == "currency-task" + assert request.context_id == "currency-context" + + +async def test_after_request_metadata_overrides_are_preserved(): + async def override_metadata(ctx, response, event): + del ctx, response + event.custom_metadata.update({ + "a2a:task_id": "caller-task", + "a2a:context_id": "caller-context", + }) + return event + + _, turns, _ = await _run_remote_task_responses( + [_bare_task_stream(_compat.TS_COMPLETED) for _ in range(3)], + config=A2aRemoteAgentConfig( + request_interceptors=[ + RequestInterceptor(after_request=override_metadata) + ] + ), + ) + for turn in turns: + assert turn[-1].custom_metadata["a2a:task_id"] == "caller-task" + assert turn[-1].custom_metadata["a2a:context_id"] == "caller-context" + assert turn[-1].custom_metadata["a2a:response"]["id"] == "currency-task" diff --git a/tests/unittests/a2a/converters/test_to_adk.py b/tests/unittests/a2a/converters/test_to_adk.py index 052271fbfe4..6eda6ff8d5a 100644 --- a/tests/unittests/a2a/converters/test_to_adk.py +++ b/tests/unittests/a2a/converters/test_to_adk.py @@ -1022,3 +1022,102 @@ def test_grounding_metadata_round_trip_message(self) -> None: back.grounding_metadata.search_entry_point.rendered_content == "test-message" ) + + +@pytest.mark.parametrize( + "state", + [ + _compat.TS_INPUT_REQUIRED, + _compat.TS_AUTH_REQUIRED, + _compat.TS_COMPLETED, + _compat.TS_FAILED, + _compat.TS_CANCELED, + _compat.TS_REJECTED, + ], +) +@pytest.mark.parametrize("empty_message", [False, True]) +@pytest.mark.parametrize("response_kind", ["task", "status", "empty-artifact"]) +def test_contentless_task_boundary_is_preserved( + state, empty_message, response_kind +): + message = ( + _compat.make_message(message_id="empty", role="agent", parts=[]) + if empty_message + else None + ) + status = _compat.make_task_status(state, message=message) + if response_kind == "status": + response = _compat.make_task_status_update_event( + task_id="task", context_id="context", status=status, final=True + ) + event = convert_a2a_status_update_to_event(response) + else: + response = _compat.make_task( + id="task", + context_id="context", + status=status, + artifacts=[_compat.make_artifact(artifact_id="empty", parts=[])] + if response_kind == "empty-artifact" + else None, + ) + event = convert_a2a_task_to_event(response) + assert event is not None + assert event.content is None + assert event.is_final_response() + assert not event.partial + + +@pytest.mark.parametrize( + "state", [_compat.TS_UNKNOWN, _compat.TS_SUBMITTED, _compat.TS_WORKING] +) +def test_contentless_progress_does_not_emit_a_final_response(state): + status = _compat.make_task_status(state) + assert ( + convert_a2a_task_to_event( + _compat.make_task(id="task", context_id="context", status=status) + ) + is None + ) + assert ( + convert_a2a_status_update_to_event( + _compat.make_task_status_update_event( + task_id="task", context_id="context", status=status, final=False + ) + ) + is None + ) + + +@pytest.mark.parametrize("response_kind", ["task", "status", "artifact"]) +def test_empty_boundary_fallback_does_not_override_part_filter(response_kind): + message = _compat.make_message( + message_id="prompt", role="agent", parts=[_compat.make_text_part("Input")] + ) + status = _compat.make_task_status(_compat.TS_INPUT_REQUIRED, message=message) + part_filter = Mock(return_value=None) + if response_kind == "status": + event = convert_a2a_status_update_to_event( + _compat.make_task_status_update_event( + task_id="task", context_id="context", status=status, final=True + ), + part_converter=part_filter, + ) + else: + event = convert_a2a_task_to_event( + _compat.make_task( + id="task", + context_id="context", + status=status, + artifacts=[ + _compat.make_artifact( + artifact_id="answer", + parts=[_compat.make_text_part("Result")], + ) + ] + if response_kind == "artifact" + else None, + ), + part_converter=part_filter, + ) + assert event is None + part_filter.assert_called()