diff --git a/docs/guides/README.md b/docs/guides/README.md index 81d8a4d5..51bc7cd9 100644 --- a/docs/guides/README.md +++ b/docs/guides/README.md @@ -8,6 +8,7 @@ This directory contains specific developer guides for the ADK Python implementat * [LlmAgent Single-Turn Mode](agents/llm_agent/single_turn.md) - Guide on using LlmAgent in single-turn mode. * [LlmAgent Task Mode](agents/llm_agent/task.md) - Guide on using LlmAgent in task mode. * [ManagedAgent](agents/managed_agent/index.md) - Guide on using ManagedAgent with server-side tools. +* [RemoteA2aAgent Task Mode](agents/remote_a2a_agent/task.md) - Guide on using RemoteA2aAgent in task mode. ### Apps * [App](apps/app/index.md) - The top-level container binding a root agent to app-wide plugins and configuration. diff --git a/docs/guides/agents/remote_a2a_agent/task.md b/docs/guides/agents/remote_a2a_agent/task.md new file mode 100644 index 00000000..a3af1cb0 --- /dev/null +++ b/docs/guides/agents/remote_a2a_agent/task.md @@ -0,0 +1,206 @@ +# RemoteA2aAgent Task Mode + +This guide explains the behavior of `RemoteA2aAgent` in `task` mode +(`mode="task"`). It covers how remote A2A agents are delegated to as sub-agents +in a multi-agent hierarchy, how the local proxy isolates task scope, and how +completion is signaled via the `finish_task` tool. + +--- + +## Introduction + +In ADK, `mode="task"` on `RemoteA2aAgent` allows a parent coordinator agent +(such as an `LlmAgent`) to delegate specific, goal-oriented sub-tasks to a +remote agent communicating over the Agent-to-Agent (A2A) protocol. + +Unlike default mode (where the remote agent acts as the primary chat interface +or a peer transfer target), a `RemoteA2aAgent` in `task` mode: + +1. **Exposed as a Tool**: The remote agent is exposed to the parent coordinator + as a tool function declaration. +2. **Session Scope Isolation**: Only conversation history relevant to the + specific sub-task execution is sent to the remote agent. +3. **Multi-Turn Interaction**: The remote agent can interact with the user + (asking clarifying questions or requesting human input) without prematurely + ending the task delegation. +4. **Completion Contract**: The task completes only when the remote agent emits + a `finish_task` tool call/response. + +--- + +## Architecture + +The following diagram illustrates how `RemoteA2aAgent` acts as a local proxy +between the parent coordinator and the remote A2A service: + +``` +┌─────────────────────────────────────────────────────────────┐ +│ Parent Agent │ +│ (e.g., LlmAgent) │ +└──────────────────────────────┬──────────────────────────────┘ + │ + │ 1. Delegates sub-task + │ via Tool Call + ▼ +┌─────────────────────────────────────────────────────────────┐ +│ RemoteA2aAgent │ +│ (Local ADK Proxy Node) │ +│ │ +│ - Reconstructs history from triggering FunctionCall.id │ +│ - Handles user interactions & pause/resume states │ +│ - Unwraps finish_task into event.output │ +│ - Propagates unrecoverable failures safely │ +└──────────────────────────────┬──────────────────────────────┘ + │ + │ 2. A2A Protocol Stream + │ (HTTP / SSE / JSON-RPC) + ▼ +┌─────────────────────────────────────────────────────────────┐ +│ Remote A2A Agent Server │ +│ (via to_a2a() Server) │ +│ │ │ +│ │ Dispatches turn │ +│ ▼ │ +│ Remote LlmAgent(mode="task") │ +│ │ +│ - Configured with mode="task" │ +│ - Automatically injects built-in `finish_task` tool │ +│ - Executes sub-task logic with local/server tools │ +│ - Calls finish_task(output=...) on completion │ +└─────────────────────────────────────────────────────────────┘ +``` + +On the remote server, configuring the underlying `LlmAgent` with `mode="task"` +causes ADK to automatically inject the `finish_task` tool and system +instructions into the model prompt. When the remote model completes its +objective, it calls `finish_task`, which the remote A2A server packages into an +A2A message for `RemoteA2aAgent` to process. + +--- + +## 1. Task Mode as a Sub-Agent + +### Behavior + +- **Tool-Based Delegation**: When attached via `sub_agents=[remote_agent]`, + the coordinator sees the remote agent's description and parameters as a + callable tool. +- **Proxy Execution**: Calling the tool suspends the parent agent and runs + `RemoteA2aAgent`. +- **History Isolation**: `RemoteA2aAgent` locates the coordinator's triggering + `FunctionCall` matching the active `isolation_scope` and scopes context to + the active task. +- **Completion Detection**: When the remote agent invokes `finish_task`, + `RemoteA2aAgent` unwraps the result into `event.output` and signals + `end_of_agent=True` to hand control back to the coordinator. + +### Example + +Here is how to define the remote A2A server and delegate to it from a parent +coordinator: + +#### Remote Server Definition (`remote_agent.py`) + +```python +from google.adk.a2a import to_a2a +from google.adk.agents import LlmAgent + +# Define the remote agent with mode="task" (automatically injects finish_task) +remote_researcher = LlmAgent( + name="researcher", + instruction="Research the given topic and call finish_task when done.", + mode="task", +) + +# Convert to an A2A server application +app = to_a2a(remote_researcher, host="localhost", port=8001) +``` + +#### Client Coordinator (`coordinator.py`) + +```python +from google.adk.agents import LlmAgent +from google.adk.agents.remote_a2a_agent import RemoteA2aAgent + +# Define the RemoteA2aAgent proxy pointing to the remote server +researcher_proxy = RemoteA2aAgent( + name="researcher", + description="Researches a topic and provides a concise summary.", + agent_card="http://localhost:8001/.well-known/agent.json", + mode="task", +) + +# Attach as a delegated sub-agent to the parent coordinator +coordinator = LlmAgent( + name="coordinator", + instruction="Write a blog post. Delegate research to the researcher agent.", + sub_agents=[researcher_proxy], +) +``` + +--- + +## 2. How it works + +### Session History Reconstruction + +When the coordinator delegates a task, `RemoteA2aAgent` scans session history to +locate the triggering `FunctionCall` matching `ctx.isolation_scope`. Only events +relevant to this specific task execution are converted into A2A messages and +sent to the remote agent. + +### User Interaction & Multi-Turn Resumption + +If the remote agent needs clarification or human input: + +1. It yields intermediate text parts or human-in-the-loop requests. +2. The framework delivers the message to the user and pauses execution. +3. When the user responds, the session resumes and routes the user's message + back to the remote agent until `finish_task` is called. + +### Failure Handling & Error Safety + +If the remote agent encounters an unrecoverable failure (`TS_FAILED`, +`TS_CANCELED`, HTTP connection failure): + +- An error event with `error_message` is yielded. +- A terminal `finish_task(result=FINISH_TASK_ERROR_RESULT)` event is generated + with `output=None` to ensure Pydantic output schemas do not fail validation + on errors. +- Control is released back to the parent coordinator cleanly with + `end_of_agent=True`. + +--- + +## 3. RemoteA2aAgent: Default Mode vs Task Mode + +This section clarifies how `mode="task"` differs from `RemoteA2aAgent`'s default +behavior: + +| Feature | Default Mode (`mode=None`) | Task Mode (`mode="task"`) | +| :--- | :--- | :--- | +| **Delegation Type** | Peer Transfer (`transfer_to_agent`) or Root Agent | Sub-Agent Tool Delegation | +| **Coordinator Exposure** | Transfer target (switches active agent) | Callable tool (`_TaskAgentTool`) | +| **History Scope** | Full session conversation history | Scoped to triggering `FunctionCall.id` | +| **Completion Mechanism** | Turn stream completion (`TS_COMPLETED`) | Explicit `finish_task` tool response | +| **Control Flow** | Control remains with remote agent until next transfer | Automatically returns control to coordinator upon task completion | +| **Output Delivery** | Streams raw text and event parts | Unwraps `finish_task` arguments into `event.output` | + +> **Note on Mode Resolution**: For `LlmAgent`, `mode=None` automatically +> resolves to `"chat"` when used as a sub-agent (making it a transfer target) or +> `"single_turn"` when used in a workflow. For `RemoteA2aAgent`, `mode=None` +> remains the default transfer target experience, while setting `mode="task"` +> explicitly enables delegated tool execution. + +--- + +## Limitations + +- **Workflow Graphs Not Supported**: `RemoteA2aAgent` in task mode + (`mode="task"`) cannot be used as a node in ADK `Workflow` graphs. It is + exclusively designed for sub-agent delegation under a parent coordinator + `LlmAgent`. +- **Requires `finish_task`**: In `mode="task"`, the remote agent must emit + `finish_task` to signal completion. +- **No Direct Transfer**: Task agents cannot be targeted via + `transfer_to_agent`; they must be invoked as sub-agents/tools. \ No newline at end of file diff --git a/src/google/adk/a2a/executor/task_result_aggregator.py b/src/google/adk/a2a/executor/task_result_aggregator.py index 42389275..60db62d2 100644 --- a/src/google/adk/a2a/executor/task_result_aggregator.py +++ b/src/google/adk/a2a/executor/task_result_aggregator.py @@ -71,6 +71,14 @@ class TaskResultAggregator: event.status.message ) event.status.state = _compat.TS_WORKING + # For backward compatibility with a2a v0.3 (a2a v1.0+ removes the `final` + # attribute from TaskStatusUpdateEvent). If we mutate the intermediate + # event to TS_WORKING, we must clear the final flag if present. + # Otherwise, the client runner will see final=True, terminate the stream + # prematurely, and miss the true final event (e.g. TS_FAILED) sent after + # the loop. + if hasattr(event, "final"): + event.final = False @property def task_state(self) -> Any: diff --git a/src/google/adk/agents/llm/task/_finish_task_tool.py b/src/google/adk/agents/llm/task/_finish_task_tool.py index 1249484f..cd747e26 100644 --- a/src/google/adk/agents/llm/task/_finish_task_tool.py +++ b/src/google/adk/agents/llm/task/_finish_task_tool.py @@ -25,7 +25,9 @@ from pydantic import TypeAdapter from pydantic import ValidationError from typing_extensions import override +from ....events.event import Event from ....tools.base_tool import BaseTool +from ....utils._schema_utils import schema_to_json_schema from ....utils._schema_utils import SchemaType from ._task_models import _DefaultTaskOutput @@ -41,6 +43,28 @@ FINISH_TASK_TOOL_NAME = 'finish_task' # passes. The wrapper uses this to distinguish a successful completion # from a validation-error retry signal. FINISH_TASK_SUCCESS_RESULT = 'Task completed.' +FINISH_TASK_ERROR_RESULT = 'Task failed.' + +# Default parameter key used to wrap primitive values for finish_task. +FINISH_TASK_DEFAULT_WRAPPER_KEY = 'result' + + +def get_output_wrapper_key( + output_schema: Optional[SchemaType], +) -> Optional[str]: + """Determine the wrapper key for the output schema.""" + schema = output_schema if output_schema is not None else _DefaultTaskOutput + if isinstance(schema, dict): + raw_schema = schema + elif isinstance(schema, types.Schema): + raw_schema = schema.model_dump(mode='json') + else: + raw_schema = schema_to_json_schema(schema) + return ( + None + if raw_schema.get('type') in ('object', 'OBJECT') + else FINISH_TASK_DEFAULT_WRAPPER_KEY + ) class FinishTaskTool(BaseTool): @@ -73,9 +97,7 @@ class FinishTaskTool(BaseTool): # FunctionDeclaration parameters must be a JSON object schema. # If the schema is already an object (e.g. BaseModel), use it directly. # Otherwise wrap it in an object with a single key. - self._wrapper_key: str | None = ( - None if raw_schema.get('type') == 'object' else 'result' - ) + self._wrapper_key: str | None = get_output_wrapper_key(raw_schema) description = ( 'Signal that this agent has completed its delegated task. Call this' @@ -182,3 +204,19 @@ no accompanying text output.""" del validated_output return FINISH_TASK_SUCCESS_RESULT + + +def is_finish_task_terminal_fr(event: Event) -> bool: + """True iff this event is a terminal FR from FinishTaskTool. + + A non-terminal FR (e.g., validation error) returns False so the + caller keeps iterating and the LLM gets a chance to retry. + """ + for fr in event.get_function_responses(): + if fr.name == FINISH_TASK_TOOL_NAME: + response = fr.response or {} + return response.get('result') in ( + FINISH_TASK_SUCCESS_RESULT, + FINISH_TASK_ERROR_RESULT, + ) + return False diff --git a/src/google/adk/agents/remote_a2a_agent.py b/src/google/adk/agents/remote_a2a_agent.py index 45785b1d..7a9a9d40 100644 --- a/src/google/adk/agents/remote_a2a_agent.py +++ b/src/google/adk/agents/remote_a2a_agent.py @@ -21,6 +21,7 @@ from pathlib import Path from typing import Any from typing import AsyncGenerator from typing import Callable +from typing import Literal from typing import Optional from typing import Union from urllib.parse import urlparse @@ -46,6 +47,7 @@ except ImportError: # Fallback for older versions of a2a-sdk. AGENT_CARD_WELL_KNOWN_PATH = "/.well-known/agent.json" + from ..a2a.agent.config import A2aRemoteAgentConfig from ..a2a.agent.interceptors.new_integration_extension import _NEW_A2A_ADK_INTEGRATION_EXTENSION from ..a2a.agent.interceptors.new_integration_extension import _new_integration_extension_interceptor @@ -66,6 +68,11 @@ from ..a2a.experimental import a2a_experimental from ..a2a.logs.log_utils import build_a2a_request_log from ..a2a.logs.log_utils import build_a2a_response_log from ..agents.invocation_context import InvocationContext +from ..agents.llm.task._finish_task_tool import FINISH_TASK_ERROR_RESULT +from ..agents.llm.task._finish_task_tool import FINISH_TASK_SUCCESS_RESULT +from ..agents.llm.task._finish_task_tool import FINISH_TASK_TOOL_NAME +from ..agents.llm.task._finish_task_tool import get_output_wrapper_key +from ..agents.llm.task._finish_task_tool import is_finish_task_terminal_fr from ..events.event import Event from ..flows.llm_flows.contents import _is_other_agent_reply from ..flows.llm_flows.contents import _present_other_agent_message @@ -73,6 +80,7 @@ from ..flows.llm_flows.functions import find_matching_function_call from ..flows.llm_flows.functions import REQUEST_CONFIRMATION_FUNCTION_CALL_NAME from ..flows.llm_flows.functions import REQUEST_EUC_FUNCTION_CALL_NAME from ..flows.llm_flows.functions import REQUEST_INPUT_FUNCTION_CALL_NAME +from ..sessions.session import Session from ..utils.context_utils import Aclosing from .base_agent import BaseAgent @@ -275,6 +283,84 @@ class A2AClientError(Exception): pass +def _text_from_content(content: Optional[genai_types.Content]) -> Optional[str]: + """Joins the text parts of a content, or None when there is no text.""" + if content is None or not content.parts: + return None + texts = [part.text for part in content.parts if part.text] + return "\n".join(texts) if texts else None + + +def _create_finish_task_event( + ctx: InvocationContext, + agent_name: str, + *, + output: Any = None, + error_message: Optional[str] = None, + is_error: bool = False, +) -> Event: + """Creates a finish_task Event.""" + return Event( + author=agent_name, + invocation_id=ctx.invocation_id, + branch=ctx.branch, + isolation_scope=ctx.isolation_scope, + error_message=error_message, + content=genai_types.Content( + role="user", + parts=[ + genai_types.Part( + function_response=genai_types.FunctionResponse( + name=FINISH_TASK_TOOL_NAME, + response={ + "result": ( + FINISH_TASK_ERROR_RESULT + if is_error + else FINISH_TASK_SUCCESS_RESULT + ) + }, + ) + ) + ], + ), + output=output, + ) + + +def _create_task_failure_events( + error_text: str, + ctx: InvocationContext, + agent_name: str, + task_id: str, + a2a_request: Any = None, +) -> tuple[Event, Event]: + """Creates events for a failed remote task.""" + error_message = f"Remote A2A task failed: {error_text}" + error_event_metadata: dict[str, Any] = { + A2A_METADATA_PREFIX + "error": error_message, + A2A_METADATA_PREFIX + "task_id": task_id, + } + if a2a_request is not None: + error_event_metadata[A2A_METADATA_PREFIX + "request"] = _compat.a2a_to_dict( + a2a_request + ) + error_event = Event( + author=agent_name, + invocation_id=ctx.invocation_id, + branch=ctx.branch, + isolation_scope=ctx.isolation_scope, + error_message=error_message, + custom_metadata=error_event_metadata, + ) + finish_event = _create_finish_task_event( + ctx=ctx, + agent_name=agent_name, + error_message=error_message, + is_error=True, + ) + return error_event, finish_event + + 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: @@ -291,6 +377,33 @@ def _add_mock_function_call(event: Event, state: TaskState) -> None: event.long_running_tool_ids = long_running_tool_ids +def _find_finish_task_args_from_history( + session: Session, + isolation_scope: Optional[str] = None, + completed_fr_event: Optional[Event] = None, +) -> Optional[dict[str, Any]]: + """Search session events for the latest finish_task FC and return args.""" + matching_fc_id = None + if completed_fr_event: + for fr in completed_fr_event.get_function_responses(): + if fr.name == FINISH_TASK_TOOL_NAME: + matching_fc_id = fr.id + break + + for event in reversed(session.events): + if isolation_scope and event.isolation_scope != isolation_scope: + continue + calls = event.get_function_calls() + for fc in calls: + if fc.name == FINISH_TASK_TOOL_NAME: + if matching_fc_id is not None: + if fc.id == matching_fc_id: + return dict(fc.args or {}) + else: + return dict(fc.args or {}) + return None + + @a2a_experimental class RemoteA2aAgent(BaseAgent): """Agent that communicates with a remote A2A agent via A2A client. @@ -307,6 +420,23 @@ class RemoteA2aAgent(BaseAgent): - Session state management across requests """ + mode: Literal["task"] | None = None + """Delegation mode. + + Only ``task`` is supported: the agent runs as a task sub-agent of a parent + ``LlmAgent`` that owns the conversation across multiple turns, then hands + control back to the parent when the remote A2A task reaches a terminal + completed state. Note: this requires the remote agent to invoke the + ``finish_task`` tool to signal completion (natively supported by ADK + task-mode agents, or must be manually implemented on custom A2A servers + by returning a FunctionResponse named ``finish_task`` with a response + containing a ``result`` key matching ``"Task completed."`` for success, or + ``"Task failed."`` for failure). Additionally, the client's ``output_schema`` + must be set to mirror the remote agent's output schema to ensure correct + output unwrapping. ``None`` (default) leaves the agent as a plain + ``transfer_to_agent`` target. + """ + def __init__( self, name: str, @@ -342,8 +472,8 @@ class RemoteA2aAgent(BaseAgent): request. full_history_when_stateless: If True, stateless agents (those that do not 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. + request. If False (default), the behavior depends on the agent's + delegation mode: True in "task" mode, False otherwise. 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. @@ -373,7 +503,7 @@ class RemoteA2aAgent(BaseAgent): self._a2a_part_converter = a2a_part_converter 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._full_history_when_stateless_param = full_history_when_stateless self._config = config or A2aRemoteAgentConfig() if not use_legacy: @@ -402,6 +532,14 @@ class RemoteA2aAgent(BaseAgent): f"got {type(agent_card)}" ) + @property + def _full_history_when_stateless(self) -> bool: + return self._full_history_when_stateless_param or self.mode == "task" + + @_full_history_when_stateless.setter + def _full_history_when_stateless(self, value: bool) -> None: + self._full_history_when_stateless_param = value + async def _ensure_httpx_client(self) -> httpx.AsyncClient: """Ensure HTTP client is available and properly configured.""" if not self._httpx_client: @@ -669,11 +807,22 @@ class RemoteA2aAgent(BaseAgent): return a2a_message def _is_remote_response(self, event: Event) -> bool: - return bool( + is_a2a_resp = bool( event.author == self.name and event.custom_metadata and event.custom_metadata.get(A2A_METADATA_PREFIX + "response", False) ) + if is_a2a_resp: + return True + + # Also stop on synthesized FR events for this agent (meaning the previous + # delegation to this agent has completed). + if self.mode == "task": + for fr in event.get_function_responses(): + if fr.name == self.name: + return True + + return False def _construct_message_parts_from_session( self, ctx: InvocationContext @@ -691,7 +840,46 @@ class RemoteA2aAgent(BaseAgent): context_id = None events_to_process = [] + task_scope = ctx.isolation_scope if self.mode == "task" else None + broke_loop = False + for event in reversed(ctx.session.events): + if task_scope: + # In task mode, we restrict the history to the current task scope + # (isolation scope) to prevent cross-task data leakage and minimize + # context size. + if event.isolation_scope == task_scope: + # Stop walking backward if we hit a previous response from this + # remote agent. Stateful remote servers already have this history + # in their session, so we don't need to resend it. + if self._is_remote_response(event): + if event.custom_metadata: + metadata = event.custom_metadata + context_id = metadata.get(A2A_METADATA_PREFIX + "context_id") + if not self._full_history_when_stateless or context_id: + broke_loop = True + break + events_to_process.append(event) + continue + # We must also include the coordinator's FunctionCall event that + # triggered this task (its ID matches the task_scope). This provides the + # remote task agent with the initial task parameters (inputs). Once we + # find it, we stop because anything older is outside this task's + # lifetime. + has_trigger_fc = False + calls = event.get_function_calls() + for fc in calls: + if fc.id == task_scope: + has_trigger_fc = True + break + if has_trigger_fc: + events_to_process.append(event) + broke_loop = True + break + # Ignore events belonging to other tasks (different isolation scopes) + # or coordinator events outside the remote task agent execution. + continue + if self._is_remote_response(event): # stop on content generated by current a2a agent given it should already # be in remote session @@ -705,9 +893,34 @@ class RemoteA2aAgent(BaseAgent): # _full_history_when_stateless is false (the default) or if the agent # is stateful (i.e. returned a context ID). if not self._full_history_when_stateless or context_id: + broke_loop = True break events_to_process.append(event) + # In task mode, an FC-delegation task must be bounded by a triggering + # FunctionCall from the coordinator. If the history walk completes to the + # root without finding the matching FC (and did not stop at a prior + # stateful turn), the isolation scope is invalid (e.g. a workflow graph + # node). + if self.mode == "task" and task_scope and not broke_loop: + raise ValueError( + f"RemoteA2aAgent '{self.name}' in task mode could not find the" + f" triggering FunctionCall for isolation scope '{task_scope}' in" + " session history. Workflow path scopes are not supported." + ) + + # Collect all FC IDs emitted by this remote agent in the task scope. + remote_fc_ids = set() + if self.mode == "task": + for event in ctx.session.events: + if ( + not task_scope or event.isolation_scope == task_scope + ) and event.author == self.name: + calls = event.get_function_calls() + for fc in calls: + if fc.id is not None: + remote_fc_ids.add(fc.id) + for event in reversed(events_to_process): processed_event: Optional[Event] = event if _is_other_agent_reply(self.name, event): @@ -731,9 +944,40 @@ class RemoteA2aAgent(BaseAgent): # path where a dropped credential resume falls back to here and the # untouched function_response would otherwise be re-serialized. continue - converted_parts = self._genai_part_converter(part) - if not isinstance(converted_parts, list): - converted_parts = [converted_parts] if converted_parts else [] + + if ( + self.mode == "task" + and task_scope + and part.function_call + and isinstance(part.function_call, genai_types.FunctionCall) + and part.function_call.id is not None + and part.function_call.id != task_scope + and part.function_call.id not in remote_fc_ids + ): + # Skip sibling function calls from the coordinator intended for other tools/agents. + continue + + if ( + self.mode == "task" + and part.function_response + and isinstance(part.function_response, genai_types.FunctionResponse) + and part.function_response.id not in remote_fc_ids + ): + # Convert non-agent function response to text to prevent A2A server + # validation errors. + text_content = ( + f"Tool {part.function_response.name} returned:" + f" {json.dumps(part.function_response.response)}" + ) + converted_parts = [_compat.make_text_part(text_content)] + else: + raw_parts = self._genai_part_converter(part) + if isinstance(raw_parts, list): + converted_parts = raw_parts + elif raw_parts is not None: + converted_parts = [raw_parts] + else: + converted_parts = [] if processed_event.author == "user": for a2a_part in converted_parts: @@ -965,152 +1209,246 @@ class RemoteA2aAgent(BaseAgent): self, ctx: InvocationContext ) -> AsyncGenerator[Event, None]: """Core implementation for async agent execution.""" + # Tracks whether task control should be released back to the parent + # coordinator and any error output to emit on early termination. + should_release_task_control = False + task_error_message: Optional[str] = None + a2a_request = None + try: - a2a_client = await self._ensure_resolved(ctx) - except Exception as e: - yield Event( - author=self.name, - error_message=f"Failed to initialize remote A2A agent: {e}", - invocation_id=ctx.invocation_id, - branch=ctx.branch, - ) - return - - # 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 - ) - - if not message_parts: - logger.warning( - "No parts to send to remote A2A agent. Emitting empty event." - ) + try: + a2a_client = await self._ensure_resolved(ctx) + except Exception as e: + task_error_message = f"Failed to initialize remote A2A agent: {e}" + should_release_task_control = True yield Event( author=self.name, - content=genai_types.Content(), + error_message=task_error_message, invocation_id=ctx.invocation_id, branch=ctx.branch, ) return - a2a_request = A2AMessage( - message_id=platform_uuid.new_uuid(), - parts=message_parts, - role=_compat.ROLE_USER, - context_id=context_id, - ) - - logger.debug(build_a2a_request_log(a2a_request)) - - try: - intercepted_request, parameters = ( - await execute_before_request_interceptors( - self._config.request_interceptors, ctx, a2a_request - ) - ) - - if isinstance(intercepted_request, Event): - yield intercepted_request - return - a2a_request = intercepted_request - - # Backward compatibility - if self._a2a_request_meta_provider: - parameters.request_metadata = self._a2a_request_meta_provider( - ctx, a2a_request + # 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 ) - # 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 - # status/artifact updates are aggregated into a running task (matching the - # 0.3.x client behavior). - normalize_stream_item = _compat.make_stream_normalizer() - async with Aclosing( - _compat.send_message( - a2a_client, - request=a2a_request, - request_metadata=parameters.request_metadata, - context=parameters.client_call_context, + if not message_parts: + logger.warning( + "No parts to send to remote A2A agent. Emitting empty event." ) - ) as agen: - async for raw_a2a_response in agen: - a2a_response = normalize_stream_item(raw_a2a_response) - logger.debug(build_a2a_response_log(a2a_response)) - - metadata = None - if isinstance(a2a_response, tuple): - task = a2a_response[0] - if task: - metadata = task.metadata - else: - metadata = a2a_response.metadata - - if metadata and _compat.metadata_get( - metadata, _NEW_A2A_ADK_INTEGRATION_EXTENSION - ): - event = await self._handle_a2a_response_v2(a2a_response, ctx) - else: - event = await self._handle_a2a_response(a2a_response, ctx) - if not event: - continue - - event = await execute_after_request_interceptors( - self._config.request_interceptors, ctx, a2a_response, event + task_error_message = "No parts to send to remote A2A agent." + should_release_task_control = True + yield Event( + author=self.name, + content=genai_types.Content(), + invocation_id=ctx.invocation_id, + branch=ctx.branch, ) - if not event: - continue + return - # Add metadata about the request and response - event.custom_metadata = event.custom_metadata or {} - event.custom_metadata[A2A_METADATA_PREFIX + "request"] = ( + a2a_request = A2AMessage( + message_id=platform_uuid.new_uuid(), + parts=message_parts, + role=_compat.ROLE_USER, + context_id=context_id, + ) + + logger.debug(build_a2a_request_log(a2a_request)) + + try: + intercepted_request, parameters = ( + await execute_before_request_interceptors( + self._config.request_interceptors, ctx, a2a_request + ) + ) + + if isinstance(intercepted_request, Event): + task_error_message = "Request intercepted" + should_release_task_control = True + yield intercepted_request + return + a2a_request = intercepted_request + + # Backward compatibility + if self._a2a_request_meta_provider: + parameters.request_metadata = self._a2a_request_meta_provider( + ctx, a2a_request + ) + + # 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 + # status/artifact updates are aggregated into a running task (matching the + # 0.3.x client behavior). + normalize_stream_item = _compat.make_stream_normalizer() + async with Aclosing( + _compat.send_message( + a2a_client, + request=a2a_request, + request_metadata=parameters.request_metadata, + context=parameters.client_call_context, + ) + ) as agen: + async for raw_a2a_response in agen: + a2a_response = normalize_stream_item(raw_a2a_response) + logger.debug(build_a2a_response_log(a2a_response)) + + task = None + metadata = None + if isinstance(a2a_response, tuple): + task = a2a_response[0] + if task: + metadata = task.metadata + else: + metadata = a2a_response.metadata + + if metadata and _compat.metadata_get( + metadata, _NEW_A2A_ADK_INTEGRATION_EXTENSION + ): + event = await self._handle_a2a_response_v2(a2a_response, ctx) + else: + event = await self._handle_a2a_response(a2a_response, ctx) + if not event: + continue + + event = await execute_after_request_interceptors( + self._config.request_interceptors, ctx, a2a_response, event + ) + if not event: + continue + + # Add metadata about the request and response + event.custom_metadata = event.custom_metadata or {} + if a2a_request: + event.custom_metadata[A2A_METADATA_PREFIX + "request"] = ( + _compat.a2a_to_dict(a2a_request) + ) + # If the response is a ClientEvent, record the task state; otherwise, + # record the message object. + if isinstance(a2a_response, tuple): + event.custom_metadata[A2A_METADATA_PREFIX + "response"] = ( + _compat.a2a_to_dict(a2a_response[0]) + ) + else: + event.custom_metadata[A2A_METADATA_PREFIX + "response"] = ( + _compat.a2a_to_dict(a2a_response) + ) + + if self.mode == "task" and is_finish_task_terminal_fr(event): + args = _find_finish_task_args_from_history( + ctx.session, ctx.isolation_scope, completed_fr_event=event + ) + if args is not None: + wrapper_key = get_output_wrapper_key(self.output_schema) + if wrapper_key and wrapper_key in args: + event.output = args[wrapper_key] + else: + event.output = args + else: + logger.warning( + "Could not find finish_task arguments in session history" + " for isolation scope '%s'. Task output will not be set.", + ctx.isolation_scope, + ) + # Yield the semantic output event so the parent runner can capture + # the final task output and record the tool response in history. + yield event + # Mark the agent as finished so parent coordinator regains control. + # Returning early terminates the stream reader, ignoring any legacy + # duplicate FRs sent by the server at the end of the run. + should_release_task_control = True + return + + yield event + + if self.mode == "task" and task: + if task.status and task.status.state in ( + _compat.TS_FAILED, + _compat.TS_CANCELED, + ): + is_cancel = task.status.state == _compat.TS_CANCELED + logger.warning( + "Remote task reported %s state. Yielding error event and " + "releasing control.", + "canceled" if is_cancel else "failure", + ) + error_text = "Unknown error" + if is_cancel: + error_text = "Task canceled" + elif event: + error_text = ( + _text_from_content(event.content) or "Unknown error" + ) + + error_event, failure_event = _create_task_failure_events( + error_text=error_text, + ctx=ctx, + agent_name=self.name, + task_id=task.id, + a2a_request=a2a_request, + ) + yield error_event + yield failure_event + should_release_task_control = True + return + + except _compat.A2A_HTTP_ERRORS as e: + error_message = f"A2A request failed: {e}" + task_error_message = error_message + should_release_task_control = True + logger.error(error_message) + status_code: object = getattr(e, "status_code", None) + custom_metadata: dict[str, Any] = { + A2A_METADATA_PREFIX + "error": error_message, + A2A_METADATA_PREFIX + "status_code": str(status_code), + } + if a2a_request: + custom_metadata[A2A_METADATA_PREFIX + "request"] = ( _compat.a2a_to_dict(a2a_request) ) - # If the response is a ClientEvent, record the task state; otherwise, - # record the message object. - if isinstance(a2a_response, tuple): - event.custom_metadata[A2A_METADATA_PREFIX + "response"] = ( - _compat.a2a_to_dict(a2a_response[0]) - ) - else: - event.custom_metadata[A2A_METADATA_PREFIX + "response"] = ( - _compat.a2a_to_dict(a2a_response) - ) + yield Event( + author=self.name, + error_message=error_message, + invocation_id=ctx.invocation_id, + branch=ctx.branch, + custom_metadata=custom_metadata, + ) - yield event + except Exception as e: + error_message = f"A2A request failed: {e}" + task_error_message = error_message + should_release_task_control = True + logger.error(error_message) + custom_metadata = { + A2A_METADATA_PREFIX + "error": error_message, + } + if a2a_request: + custom_metadata[A2A_METADATA_PREFIX + "request"] = ( + _compat.a2a_to_dict(a2a_request) + ) + yield Event( + author=self.name, + error_message=error_message, + invocation_id=ctx.invocation_id, + branch=ctx.branch, + custom_metadata=custom_metadata, + ) - except _compat.A2A_HTTP_ERRORS as e: - error_message = f"A2A request failed: {e}" - logger.error(error_message) - status_code: object = getattr(e, "status_code", None) - yield Event( - author=self.name, - error_message=error_message, - invocation_id=ctx.invocation_id, - branch=ctx.branch, - custom_metadata={ - A2A_METADATA_PREFIX + "request": _compat.a2a_to_dict(a2a_request), - A2A_METADATA_PREFIX + "error": error_message, - A2A_METADATA_PREFIX + "status_code": str(status_code), - }, - ) - - except Exception as e: - error_message = f"A2A request failed: {e}" - logger.error(error_message) - - yield Event( - author=self.name, - error_message=error_message, - invocation_id=ctx.invocation_id, - branch=ctx.branch, - custom_metadata={ - A2A_METADATA_PREFIX + "request": _compat.a2a_to_dict(a2a_request), - A2A_METADATA_PREFIX + "error": error_message, - }, - ) + finally: + if self.mode == "task" and should_release_task_control: + if task_error_message is not None: + yield _create_finish_task_event( + ctx=ctx, + agent_name=self.name, + error_message=task_error_message, + is_error=True, + ) + ctx.set_agent_state(self.name, end_of_agent=True) + yield self._create_agent_state_event(ctx) async def _run_live_impl( self, ctx: InvocationContext @@ -1152,8 +1490,10 @@ class RemoteA2aAgent(BaseAgent): """ promoted = False async for event in super()._run_impl(ctx=ctx, node_input=node_input): - if not promoted and self._promote_response_to_output( - event, ctx.node_path + if ( + self.mode != "task" + and not promoted + and self._promote_response_to_output(event, ctx.node_path) ): promoted = True yield event diff --git a/src/google/adk/runners.py b/src/google/adk/runners.py index 1d3965a4..3a635887 100644 --- a/src/google/adk/runners.py +++ b/src/google/adk/runners.py @@ -40,6 +40,7 @@ from .agents.context_cache_config import ContextCacheConfig from .agents.invocation_context import InvocationContext from .agents.invocation_context import new_invocation_context_id from .agents.live_request_queue import LiveRequestQueue +from .agents.llm.task._finish_task_tool import FINISH_TASK_ERROR_RESULT from .agents.llm.task._finish_task_tool import FINISH_TASK_SUCCESS_RESULT from .agents.llm.task._finish_task_tool import FINISH_TASK_TOOL_NAME from .agents.run_config import RunConfig @@ -115,12 +116,13 @@ def _find_active_task_scope(session: Session) -> Optional[tuple[str, str]]: scope = ``@``, stamped on every event the task agent emits. - Both close on a SUCCESSFUL ``finish_task`` FunctionResponse — - i.e., one whose response is ``FINISH_TASK_SUCCESS_RESULT``. An - error FR (validation failure) does NOT close the scope: the task - agent is still active, will see the error, and retry. Walking - backward, the first non-empty scope we encounter that hasn't been - closed by a later successful ``finish_task`` is the paused task + Both close on a terminal ``finish_task`` FunctionResponse containing a + 'result' key matching ``FINISH_TASK_SUCCESS_RESULT`` or + ``FINISH_TASK_ERROR_RESULT``. A FunctionResponse containing an 'error' key + (indicating a tool validation failure) does NOT close the scope: the task + agent is still active, will see the validation error, and retry. Walking + backward, the first non-empty scope we encounter that hasn't been closed by a + later successful or failed terminal ``finish_task`` is the paused task awaiting the user's next reply. Used by ``Runner._append_user_event`` to scope the new user message @@ -130,8 +132,12 @@ def _find_active_task_scope(session: Session) -> Optional[tuple[str, str]]: A tuple of (isolation_scope, invocation_id) for the active task if found, or None if no active task scope is found. """ + # Pass 1: Scan forward to find all scopes that have successfully finished. + # We must do this in a separate pass because walking backward directly would + # hit post-finish events (like status updates or duplicate FRs) before hitting + # the older success FR, falsely indicating the scope is still active. finished_scopes: set[str] = set() - for event in reversed(session.events): + for event in session.events: scope = event.isolation_scope if not scope: continue @@ -140,9 +146,18 @@ def _find_active_task_scope(session: Session) -> Optional[tuple[str, str]]: fr = part.function_response if fr and fr.name == FINISH_TASK_TOOL_NAME: response = fr.response or {} - if response.get('result') == FINISH_TASK_SUCCESS_RESULT: + if response.get('result') in ( + FINISH_TASK_SUCCESS_RESULT, + FINISH_TASK_ERROR_RESULT, + ): finished_scopes.add(scope) break + + # Pass 2: Walk backward to find the latest active scope that is not finished. + for event in reversed(session.events): + scope = event.isolation_scope + if not scope: + continue if scope not in finished_scopes: return scope, event.invocation_id return None @@ -1203,6 +1218,15 @@ class Runner: from .agents.llm_agent import LlmAgent from .workflow._base_node import BaseNode + # Optional dependency: RemoteA2aAgent is only available if a2a is installed. + remote_a2a_agent_type: Any = None + try: + from .agents.remote_a2a_agent import RemoteA2aAgent # pylint: disable=g-import-not-at-top + + remote_a2a_agent_type = RemoteA2aAgent + except ImportError: + pass + if isinstance(self.agent, LlmAgent): if self.agent.mode is None: # LlmAgent as root agent defaults to chat mode. @@ -1223,8 +1247,14 @@ class Runner: # when the chat coordinator has task-mode sub-agents, # the wrapper handles delegation via ctx.run_node. Don't let # the legacy sub-agent picker bypass the coordinator on resume. + remote_a2a_agent_class = ( + (remote_a2a_agent_type,) + if remote_a2a_agent_type is not None + else () + ) has_task_subagent = any( - isinstance(sa, LlmAgent) and getattr(sa, 'mode', None) == 'task' + isinstance(sa, (LlmAgent,) + remote_a2a_agent_class) + and getattr(sa, 'mode', None) == 'task' for sa in self.agent.sub_agents or [] ) agent_to_run: BaseAgent diff --git a/src/google/adk/workflow/_llm_agent_wrapper.py b/src/google/adk/workflow/_llm_agent_wrapper.py index f1209781..4ef60d34 100644 --- a/src/google/adk/workflow/_llm_agent_wrapper.py +++ b/src/google/adk/workflow/_llm_agent_wrapper.py @@ -26,8 +26,8 @@ from typing import TYPE_CHECKING from google.genai import types from ..agents.context import Context -from ..agents.llm.task._finish_task_tool import FINISH_TASK_SUCCESS_RESULT from ..agents.llm.task._finish_task_tool import FINISH_TASK_TOOL_NAME as _FINISH_TASK_FC_NAME +from ..agents.llm.task._finish_task_tool import is_finish_task_terminal_fr from ..events.event import Event from ..flows.llm_flows.functions import REQUEST_CONFIRMATION_FUNCTION_CALL_NAME from ..utils._schema_utils import validate_schema @@ -47,19 +47,6 @@ def _extract_finish_task_fc(event: Event) -> types.FunctionCall | None: return None -def _is_finish_task_success_fr(event: Event) -> bool: - """True iff this event is the success FR from FinishTaskTool. - - A non-success FR (e.g., validation error) returns False so the - caller keeps iterating and the LLM gets a chance to retry. - """ - for fr in event.get_function_responses(): - if fr.name == _FINISH_TASK_FC_NAME: - response = fr.response or {} - return response.get('result') == FINISH_TASK_SUCCESS_RESULT - return False - - def _extract_task_delegation_fcs( event: Event, tools_dict: Mapping[str, ToolUnion] ) -> list[types.FunctionCall]: @@ -532,7 +519,7 @@ async def run_llm_agent_as_node( yield event continue - if pending_fc_args is not None and _is_finish_task_success_fr(event): + if pending_fc_args is not None and is_finish_task_terminal_fr(event): wrapper_key = getattr(finish_tool, '_wrapper_key', None) if wrapper_key and wrapper_key in pending_fc_args: event.output = pending_fc_args[wrapper_key] diff --git a/src/google/adk/workflow/utils/_workflow_graph_utils.py b/src/google/adk/workflow/utils/_workflow_graph_utils.py index b35c78dc..57f13dbc 100644 --- a/src/google/adk/workflow/utils/_workflow_graph_utils.py +++ b/src/google/adk/workflow/utils/_workflow_graph_utils.py @@ -17,6 +17,7 @@ from __future__ import annotations from typing import Any +from typing import cast from typing import Literal from ...tools.base_tool import BaseTool @@ -79,6 +80,15 @@ def build_node( # workflow_graph_utils -> agents.llm_agent -> ... -> workflow_graph_utils from ...agents.llm_agent import LlmAgent + # Optional dependency: RemoteA2aAgent is only available if a2a is installed. + _remote_a2a_agent_type: Any = None + try: + from ...agents.remote_a2a_agent import RemoteA2aAgent # pylint: disable=g-import-not-at-top + + _remote_a2a_agent_type = RemoteA2aAgent + except ImportError: + pass + if isinstance(node_like, BaseNode): kwargs: dict[str, Any] = {} if name is not None: @@ -90,14 +100,27 @@ def build_node( if timeout is not None: kwargs['timeout'] = timeout - if isinstance(node_like, LlmAgent): + is_remote_a2a_task = False + if _remote_a2a_agent_type is not None: + is_remote_a2a_task = ( + isinstance(node_like, _remote_a2a_agent_type) + and node_like.mode == 'task' + ) + if is_remote_a2a_task and getattr(node_like, 'parent_agent', None) is None: + raise ValueError( + 'RemoteA2aAgent in task mode is not supported as a standalone ' + 'workflow node. It is only supported in tool-delegation mode.' + ) + + if isinstance(node_like, LlmAgent) or is_remote_a2a_task: if rerun_on_resume is None: kwargs['rerun_on_resume'] = True - agent = node_like.clone(update=kwargs) + agent_node = cast(Any, node_like) + agent = agent_node.clone(update=kwargs) # Preserve parent agent reference that was lost during clone - agent.parent_agent = node_like.parent_agent + agent.parent_agent = agent_node.parent_agent - if agent.mode is None: + if isinstance(agent, LlmAgent) and agent.mode is None: # Sub-agents dynamically attached to a parent agent default to 'chat' # mode to enable agent transfer. # Standalone agents in a workflow graph default to 'single_turn'. @@ -109,12 +132,12 @@ def build_node( if agent.mode in ('task', 'chat'): agent.wait_for_output = True - if agent.parallel_worker: + if isinstance(agent, LlmAgent) and agent.parallel_worker: from .._parallel_worker import _ParallelWorker agent.parallel_worker = False return _ParallelWorker(node=agent) - return agent + return cast(BaseNode, agent) else: if kwargs: return node_like.model_copy(update=kwargs) diff --git a/tests/unittests/a2a/executor/test_task_result_aggregator.py b/tests/unittests/a2a/executor/test_task_result_aggregator.py index bca5b4e5..18a253e7 100644 --- a/tests/unittests/a2a/executor/test_task_result_aggregator.py +++ b/tests/unittests/a2a/executor/test_task_result_aggregator.py @@ -329,3 +329,25 @@ class TestTaskResultAggregator: assert ( self.aggregator.task_status_message == auth_message ) # Message unchanged because task state is not working + + def test_process_final_non_working_event_mutates_final_and_state(self): + """Test that a final=True non-working event is mutated to final=False and state=TS_WORKING.""" + status_message = create_test_message("Failed") + event = _compat.make_task_status_update_event( + task_id="test-task", + context_id="test-context", + status=_compat.make_task_status( + _compat.TS_FAILED, message=status_message + ), + final=True, + ) + + self.aggregator.process_event(event) + + # Aggregator should record the true state + assert self.aggregator.task_state == _compat.TS_FAILED + + # Event itself must be mutated for the wire protocol + assert event.status.state == _compat.TS_WORKING + if hasattr(event, "final"): + assert event.final is False diff --git a/tests/unittests/agents/test_remote_a2a_agent.py b/tests/unittests/agents/test_remote_a2a_agent.py index a4568835..1fb8002e 100644 --- a/tests/unittests/agents/test_remote_a2a_agent.py +++ b/tests/unittests/agents/test_remote_a2a_agent.py @@ -44,6 +44,9 @@ from google.adk.a2a.agent.utils import execute_after_request_interceptors from google.adk.a2a.agent.utils import execute_before_card_request_interceptors from google.adk.a2a.agent.utils import execute_before_request_interceptors from google.adk.agents.invocation_context import InvocationContext +from google.adk.agents.llm.task._finish_task_tool import FINISH_TASK_ERROR_RESULT +from google.adk.agents.llm.task._finish_task_tool import FINISH_TASK_SUCCESS_RESULT +from google.adk.agents.llm.task._finish_task_tool import FINISH_TASK_TOOL_NAME from google.adk.agents.remote_a2a_agent import A2A_METADATA_PREFIX from google.adk.agents.remote_a2a_agent import AgentCardResolutionError from google.adk.agents.remote_a2a_agent import RemoteA2aAgent @@ -52,6 +55,7 @@ from google.adk.events.event import Event from google.adk.sessions.session import Session from google.genai import types as genai_types import httpx +from pydantic import BaseModel import pytest @@ -110,6 +114,17 @@ def _make_stream_message(message: A2AMessage): return message +def _make_stream_task(task: A2ATask): + """Wrap a Task in the shape ``send_message`` yields for the active SDK.""" + if _compat.IS_A2A_V1: + from a2a.types import StreamResponse + + resp = StreamResponse() + resp.task.CopyFrom(task) + return resp + return (task, None) + + def _make_artifact_chunk(text: str, *, append: bool, last_chunk: bool): """Build one streamed chunk of an artifact, version-agnostically.""" return TaskArtifactUpdateEvent( @@ -143,6 +158,21 @@ def _make_accumulated_task(part_texts): ) +def _make_dummy_task_trigger_event( + task_id: str = "task-1", agent_name: str = "test_agent" +) -> Event: + """Build a dummy triggering Event containing a FunctionCall for task delegation.""" + trigger_fc = genai_types.FunctionCall( + id=task_id, name=agent_name, args={"request": "start"} + ) + return Event( + author="coordinator", + content=genai_types.Content( + role="model", parts=[genai_types.Part(function_call=trigger_fc)] + ), + ) + + # Helper function to create a proper AgentCard for testing def create_test_agent_card( name: str = "test-agent", @@ -1252,43 +1282,102 @@ class TestRemoteA2aAgentMessageHandling: assert parts == [] assert context_id is None - def test_construct_message_parts_from_session_stops_on_agent_reply(self): - """Test message parts construction stops on agent reply by default.""" + def test_construct_message_parts_from_session_foreign_function_response_not_converted( + self, + ): + """Test that foreign function responses are NOT converted to text in default mode.""" + # Mock event with a function response + mock_fr = genai_types.FunctionResponse( + id="fc-1", name="tool_1", response={"result": "done"} + ) + mock_part = Mock() + mock_part.function_response = mock_fr + mock_part.text = None + + mock_content = Mock() + mock_content.parts = [mock_part] + + mock_event = Mock() + mock_event.author = "user" + mock_event.content = mock_content + mock_event.get_function_calls.return_value = [] + mock_event.get_function_responses.return_value = [] + + self.mock_session.events = [mock_event] + + with patch( + "google.adk.agents.remote_a2a_agent._present_other_agent_message" + ) as mock_present: + mock_present.return_value = mock_event + + mock_a2a_part = _compat.make_text_part("tool_response_text") + self.mock_genai_part_converter.return_value = mock_a2a_part + + parts, _ = self.agent._construct_message_parts_from_session( + self.mock_context + ) + + # Should call the converter, not convert to text + self.mock_genai_part_converter.assert_called_once_with(mock_part) + assert len(parts) == 1 + assert parts[0] == mock_a2a_part + + def test_construct_message_parts_from_session_stops_on_agent_reply_when_disabled( + self, + ): + """Test message parts construction stops on agent reply when disabled.""" + self.agent._full_history_when_stateless = False part1 = Mock() part1.text = "User 1" content1 = Mock() content1.parts = [part1] - user1 = Mock() - user1.content = content1 - user1.author = "user" - user1.custom_metadata = None + user1 = Mock( + live_session_id=None, + author="user", + custom_metadata=None, + content=content1, + ) + user1.get_function_calls.return_value = [] + user1.get_function_responses.return_value = [] part2 = Mock() part2.text = "Agent 1" content2 = Mock() content2.parts = [part2] - agent1 = Mock() - agent1.content = content2 - agent1.author = self.agent.name - agent1.custom_metadata = { - A2A_METADATA_PREFIX + "response": True, - } + agent1 = Mock( + live_session_id=None, + author=self.agent.name, + content=content2, + custom_metadata={ + A2A_METADATA_PREFIX + "response": True, + }, + ) + agent1.get_function_calls.return_value = [] + agent1.get_function_responses.return_value = [] - agent2 = Mock() - agent2.content = None - agent2.author = self.agent.name - # Just actions, no content. Not marked as a response. - agent2.actions = Mock() - agent2.custom_metadata = None + agent2 = Mock( + live_session_id=None, + author=self.agent.name, + content=None, + # Just actions, no content. Not marked as a response. + actions=Mock(), + custom_metadata=None, + ) + agent2.get_function_calls.return_value = [] + agent2.get_function_responses.return_value = [] part3 = Mock() part3.text = "User 2" content3 = Mock() content3.parts = [part3] - user2 = Mock() - user2.content = content3 - user2.author = "user" - user2.custom_metadata = None + user2 = Mock( + live_session_id=None, + author="user", + content=content3, + custom_metadata=None, + ) + user2.get_function_calls.return_value = [] + user2.get_function_responses.return_value = [] self.mock_session.events = [user1, agent1, user2, agent2] @@ -1308,35 +1397,51 @@ class TestRemoteA2aAgentMessageHandling: assert _compat.part_text(parts[0]) == "User 2" assert context_id is None - def test_construct_message_parts_from_session_stateless_full_history(self): + def test_construct_message_parts_from_session_stateless_full_history_when_enabled( + self, + ): """Test full history for stateless agent when enabled.""" self.agent._full_history_when_stateless = True part1 = Mock() part1.text = "User 1" content1 = Mock() content1.parts = [part1] - user1 = Mock() - user1.content = content1 - user1.author = "user" - user1.custom_metadata = None + user1 = Mock( + live_session_id=None, + author="user", + custom_metadata=None, + content=content1, + ) + user1.get_function_calls.return_value = [] + user1.get_function_responses.return_value = [] part2 = Mock() part2.text = "Agent 1" content2 = Mock() content2.parts = [part2] - agent1 = Mock() - agent1.content = content2 - agent1.author = self.agent.name - agent1.custom_metadata = None + agent1 = Mock( + live_session_id=None, + author=self.agent.name, + content=content2, + custom_metadata={ + A2A_METADATA_PREFIX + "response": True, + }, + ) + agent1.get_function_calls.return_value = [] + agent1.get_function_responses.return_value = [] part3 = Mock() part3.text = "User 2" content3 = Mock() content3.parts = [part3] - user2 = Mock() - user2.content = content3 - user2.author = "user" - user2.custom_metadata = None + user2 = Mock( + live_session_id=None, + author="user", + content=content3, + custom_metadata=None, + ) + user2.get_function_calls.return_value = [] + user2.get_function_responses.return_value = [] self.mock_session.events = [user1, agent1, user2] @@ -2078,6 +2183,480 @@ class TestRemoteA2aAgentMessageHandling: assert result is None +class TestRemoteA2aAgentTaskModeMessageHandling: + """Test message handling functionality under task mode.""" + + def setup_method(self): + """Setup test fixtures.""" + self.agent_card = create_test_agent_card() + self.mock_genai_part_converter = Mock() + self.mock_a2a_part_converter = Mock() + self.agent = RemoteA2aAgent( + name="test_agent", + agent_card=self.agent_card, + genai_part_converter=self.mock_genai_part_converter, + a2a_part_converter=self.mock_a2a_part_converter, + mode="task", + ) + + # Mock session and context + self.mock_session = Mock(spec=Session) + self.mock_session.id = "session-123" + self.mock_session.events = [] + + self.mock_context = Mock(spec=InvocationContext) + self.mock_context.session = self.mock_session + self.mock_context.invocation_id = "invocation-123" + self.mock_context.branch = "main" + self.mock_context.isolation_scope = "task-1" + + def test_construct_message_parts_from_session_isolates_history(self): + """Test history collection in task mode isolates to current scope.""" + # 1. Event outside task (oldest) + event_outside_old = Mock( + live_session_id=None, + isolation_scope=None, + author="user", + ) + event_outside_old.get_function_calls.return_value = [] + event_outside_old.get_function_responses.return_value = [] + + # 2. Trigger FC event (scope=None, but contains FC with id="task-1") + trigger_fc = genai_types.FunctionCall( + id="task-1", name="test_agent", args={"request": "start"} + ) + trigger_part = Mock() + trigger_part.text = "Trigger message" + trigger_part.function_response = None + + trigger_event = Mock( + live_session_id=None, + isolation_scope=None, + author="deal_agent", + ) + trigger_event.get_function_calls.return_value = [trigger_fc] + trigger_event.get_function_responses.return_value = [] + trigger_event.content = Mock() + trigger_event.content.parts = [trigger_part] + + # 3. Event inside task (remote response) + remote_part = Mock() + remote_part.text = "remote task agent output" + remote_part.function_response = None + + event_inside_remote = Mock( + live_session_id=None, + isolation_scope="task-1", + author=self.agent.name, + custom_metadata={ + A2A_METADATA_PREFIX + "response": True, + A2A_METADATA_PREFIX + "context_id": "ctx-1", + }, + ) + event_inside_remote.get_function_calls.return_value = [] + event_inside_remote.get_function_responses.return_value = [] + event_inside_remote.content = Mock() + event_inside_remote.content.parts = [remote_part] + + # 4. Event inside task (user reply) + user_part = Mock() + user_part.text = "User reply" + user_part.function_response = None + + event_inside_user = Mock( + live_session_id=None, + isolation_scope="task-1", + author="user", + ) + event_inside_user.get_function_calls.return_value = [] + event_inside_user.get_function_responses.return_value = [] + event_inside_user.content = Mock() + event_inside_user.content.parts = [user_part] + + # 5. Event outside task (newer) + event_outside_new = Mock( + live_session_id=None, + isolation_scope="task-2", + author="other_agent", + ) + event_outside_new.get_function_calls.return_value = [] + event_outside_new.get_function_responses.return_value = [] + + # Session events (oldest to newest) + self.mock_session.events = [ + event_outside_old, + trigger_event, + event_inside_remote, + event_inside_user, + event_outside_new, + ] + + # Mock converter to return a real text Part + def mock_converter(part): + return _compat.make_text_part(getattr(part, "text", "default")) + + self.mock_genai_part_converter.side_effect = mock_converter + + with patch( + "google.adk.agents.remote_a2a_agent._present_other_agent_message" + ) as mock_present: + mock_present.side_effect = lambda event: event + + parts, context_id = self.agent._construct_message_parts_from_session( + self.mock_context + ) + + # Stateful resumption: should stop at event_inside_remote and only collect + # event_inside_user + assert len(parts) == 1 + assert _compat.part_text(parts[0]) == "User reply" + assert context_id == "ctx-1" + + def test_construct_message_parts_from_session_first_turn(self): + """Test history collection in task mode first turn (collects up to trigger FC).""" + # 1. Event outside task (oldest) + event_outside_old = Mock( + live_session_id=None, + isolation_scope=None, + author="user", + ) + event_outside_old.get_function_calls.return_value = [] + event_outside_old.get_function_responses.return_value = [] + + # 2. Trigger FC event + trigger_fc = genai_types.FunctionCall( + id="task-1", name="test_agent", args={"request": "start"} + ) + trigger_part = Mock() + trigger_part.text = "Trigger message" + trigger_part.function_response = None + + trigger_event = Mock( + live_session_id=None, + isolation_scope=None, + author="deal_agent", + ) + trigger_event.get_function_calls.return_value = [trigger_fc] + trigger_event.get_function_responses.return_value = [] + trigger_event.content = Mock() + trigger_event.content.parts = [trigger_part] + + # Session events (oldest to newest) + self.mock_session.events = [ + event_outside_old, + trigger_event, + ] + + def mock_converter(part): + return _compat.make_text_part(getattr(part, "text", "default")) + + self.mock_genai_part_converter.side_effect = mock_converter + + with patch( + "google.adk.agents.remote_a2a_agent._present_other_agent_message" + ) as mock_present: + mock_present.side_effect = lambda event: event + + parts, context_id = self.agent._construct_message_parts_from_session( + self.mock_context + ) + + # First turn: collects only trigger_event + assert len(parts) == 1 + assert _compat.part_text(parts[0]) == "Trigger message" + assert context_id is None + + def test_construct_message_parts_from_session_stateless_full_history(self): + """Test full history for stateless agent in task mode when enabled.""" + self.agent._full_history_when_stateless = True + + # 1. Event outside task (oldest) + event_outside_old = Mock( + live_session_id=None, + isolation_scope=None, + author="user", + ) + event_outside_old.get_function_calls.return_value = [] + event_outside_old.get_function_responses.return_value = [] + + # 2. Trigger FC event + trigger_fc = genai_types.FunctionCall( + id="task-1", name="test_agent", args={"request": "start"} + ) + trigger_part = Mock() + trigger_part.text = "Trigger message" + trigger_part.function_response = None + + trigger_event = Mock( + live_session_id=None, + isolation_scope=None, + author="deal_agent", + ) + trigger_event.get_function_calls.return_value = [trigger_fc] + trigger_event.get_function_responses.return_value = [] + trigger_event.content = Mock() + trigger_event.content.parts = [trigger_part] + + # 3. Event inside task (remote response) - STATELESS (no context_id) + remote_part = Mock() + remote_part.text = "remote task agent output" + remote_part.function_response = None + + event_inside_remote = Mock( + live_session_id=None, + isolation_scope="task-1", + author=self.agent.name, + custom_metadata={ + A2A_METADATA_PREFIX + "response": True, + # NO context_id + }, + ) + event_inside_remote.get_function_calls.return_value = [] + event_inside_remote.get_function_responses.return_value = [] + event_inside_remote.content = Mock() + event_inside_remote.content.parts = [remote_part] + + # 4. Event inside task (user reply) + user_part = Mock() + user_part.text = "User reply" + user_part.function_response = None + + event_inside_user = Mock( + live_session_id=None, + isolation_scope="task-1", + author="user", + ) + event_inside_user.get_function_calls.return_value = [] + event_inside_user.get_function_responses.return_value = [] + event_inside_user.content = Mock() + event_inside_user.content.parts = [user_part] + + # Session events (oldest to newest) + self.mock_session.events = [ + event_outside_old, + trigger_event, + event_inside_remote, + event_inside_user, + ] + + def mock_converter(part): + return _compat.make_text_part(getattr(part, "text", "default")) + + self.mock_genai_part_converter.side_effect = mock_converter + + with patch( + "google.adk.agents.remote_a2a_agent._present_other_agent_message" + ) as mock_present: + mock_present.side_effect = lambda event: event + + parts, context_id = self.agent._construct_message_parts_from_session( + self.mock_context + ) + + # Stateless resumption with full history enabled: + # Should NOT stop at event_inside_remote. + # Should collect: event_inside_user, event_inside_remote, trigger_event. + assert len(parts) == 3 + assert _compat.part_text(parts[0]) == "Trigger message" + assert _compat.part_text(parts[1]) == "remote task agent output" + assert _compat.part_text(parts[2]) == "User reply" + assert context_id is None + + def test_construct_message_parts_from_session_filters_sibling_fcs(self): + """Test that sibling FunctionCalls from the coordinator are filtered out.""" + trigger_fc = genai_types.FunctionCall( + id="task-1", name="test_agent", args={"request": "start"} + ) + sibling_fc = genai_types.FunctionCall( + id="sibling-2", name="other_tool", args={"other": "data"} + ) + + trigger_part = Mock() + trigger_part.function_call = trigger_fc + trigger_part.function_response = None + trigger_part.text = None + + sibling_part = Mock() + sibling_part.function_call = sibling_fc + sibling_part.function_response = None + sibling_part.text = None + + trigger_event = Mock( + live_session_id=None, + isolation_scope=None, + author="coordinator", + ) + trigger_event.get_function_calls.return_value = [trigger_fc, sibling_fc] + trigger_event.get_function_responses.return_value = [] + trigger_event.content = Mock() + trigger_event.content.parts = [trigger_part, sibling_part] + + self.mock_session.events = [trigger_event] + + def mock_converter(part): + return _compat.make_text_part(f"FC:{part.function_call.id}") + + self.mock_genai_part_converter.side_effect = mock_converter + + with patch( + "google.adk.agents.remote_a2a_agent._present_other_agent_message" + ) as mock_present: + mock_present.side_effect = lambda event: event + + parts, context_id = self.agent._construct_message_parts_from_session( + self.mock_context + ) + + assert len(parts) == 1 + assert _compat.part_text(parts[0]) == "FC:task-1" + assert context_id is None + + def test_construct_message_parts_from_session_foreign_function_response_converted( + self, + ): + """Test that foreign function responses ARE converted to text in task mode.""" + # Mock event with a function response + mock_fr = genai_types.FunctionResponse( + id="fc-1", name="tool_1", response={"result": "done"} + ) + mock_part = Mock() + mock_part.function_response = mock_fr + mock_part.text = None + + mock_content = Mock() + mock_content.parts = [mock_part] + + mock_event = Mock() + mock_event.isolation_scope = "task-1" + mock_event.author = "user" + mock_event.content = mock_content + mock_event.get_function_calls.return_value = [] + mock_event.get_function_responses.return_value = [] + + # Trigger event + trigger_fc = genai_types.FunctionCall( + id="task-1", name="test_agent", args={"request": "start"} + ) + trigger_part = Mock() + trigger_part.text = "Trigger message" + trigger_part.function_response = None + trigger_event = Mock( + live_session_id=None, + isolation_scope=None, + author="deal_agent", + ) + trigger_event.get_function_calls.return_value = [trigger_fc] + trigger_event.get_function_responses.return_value = [] + trigger_event.content = Mock() + trigger_event.content.parts = [trigger_part] + + self.mock_session.events = [trigger_event, mock_event] + + # Setup converter to return distinguishable parts + def mock_converter(part): + return _compat.make_text_part(getattr(part, "text", "default")) + + self.mock_genai_part_converter.side_effect = mock_converter + + with patch( + "google.adk.agents.remote_a2a_agent._present_other_agent_message" + ) as mock_present: + mock_present.side_effect = lambda event: event + + parts, _ = self.agent._construct_message_parts_from_session( + self.mock_context + ) + + assert len(parts) == 2 + assert _compat.part_text(parts[0]) == "Trigger message" + # The foreign FR should be converted to text + expected_text = 'Tool tool_1 returned: {"result": "done"}' + assert _compat.part_text(parts[1]) == expected_text + # Check that converter was not called for the foreign FR + self.mock_genai_part_converter.assert_called_once_with(trigger_part) + + def test_construct_message_parts_from_session_non_foreign_fr_not_converted_when_fc_before_break( + self, + ): + """Test that non-foreign FR is NOT converted to text even if its FC was before context_id break.""" + # 1. Trigger event (matching task-1) + trigger_fc = genai_types.FunctionCall( + id="task-1", name="test_agent", args={"request": "start"} + ) + trigger_part = Mock() + trigger_part.text = "Trigger message" + trigger_part.function_response = None + trigger_event = Mock( + live_session_id=None, isolation_scope=None, author="deal_agent" + ) + trigger_event.get_function_calls.return_value = [trigger_fc] + trigger_event.get_function_responses.return_value = [] + trigger_event.content = Mock() + trigger_event.content.parts = [trigger_part] + + # 2. Remote agent event (FC: input request) - this will be before the break + remote_fc = genai_types.FunctionCall( + id="fc-input-req", + name="user_input_tool", + args={"prompt": "enter value"}, + ) + remote_part = Mock() + remote_part.function_call = remote_fc + remote_part.text = "Please enter value" + remote_event = Mock( + live_session_id=None, isolation_scope="task-1", author="test_agent" + ) + remote_event.get_function_calls.return_value = [remote_fc] + remote_event.get_function_responses.return_value = [] + remote_event.content = Mock() + remote_event.content.parts = [remote_part] + # This event has context_id, so it will trigger the break in history walk + remote_event.metadata = {A2A_METADATA_PREFIX + "context_id": "context-old"} + + # 3. User event (FR: user input response) - this will be after the break + user_fr = genai_types.FunctionResponse( + id="fc-input-req", name="user_input_tool", response={"result": "my-val"} + ) + user_part = Mock() + user_part.function_response = user_fr + user_part.text = None + user_event = Mock( + live_session_id=None, isolation_scope="task-1", author="user" + ) + user_event.get_function_calls.return_value = [] + user_event.get_function_responses.return_value = [user_fr] + user_event.content = Mock() + user_event.content.parts = [user_part] + + # Session events order (chronological): trigger, remote (break), user + self.mock_session.events = [trigger_event, remote_event, user_event] + + # Setup converter to return distinguishable parts + mock_a2a_part = Mock() + self.mock_genai_part_converter.return_value = [mock_a2a_part] + + with ( + patch( + "google.adk.agents.remote_a2a_agent._present_other_agent_message" + ) as mock_present, + patch( + "google.adk.agents.remote_a2a_agent._compat.part_metadata" + ) as mock_part_metadata, + ): + mock_present.side_effect = lambda event: event + mock_part_metadata.return_value = {} + + parts, _ = self.agent._construct_message_parts_from_session( + self.mock_context + ) + + # We expect the FR to NOT be converted to text, so the converter is called + # and we get the mock_a2a_part back. + assert len(parts) == 1 + assert parts[0] == mock_a2a_part + self.mock_genai_part_converter.assert_called_once_with(user_part) + + class TestRemoteA2aAgentStreamingArtifactChunks: """Regression tests for chunked artifact streams.""" @@ -3248,6 +3827,192 @@ class TestRemoteA2aAgentExecution: assert len(events) == 1 assert "A2A request failed" in events[0].error_message + @pytest.mark.asyncio + async def test_run_async_impl_task_mode_rejects_missing_trigger_fc(self): + """Test _run_async_impl raises ValueError when isolation_scope has no matching trigger FC.""" + agent = RemoteA2aAgent( + name="test_agent", + agent_card=self.agent_card, + mode="task", + ) + self.mock_context.isolation_scope = "workflow_path/node@run1" + self.mock_session.events = [] + + with pytest.raises( + ValueError, match="could not find the triggering FunctionCall" + ): + _ = [e async for e in agent._run_async_impl(self.mock_context)] + + @pytest.mark.asyncio + async def test_run_async_impl_task_mode_releases_control_on_init_failure( + self, + ): + """Test _run_async_impl in task mode releases control on initialization failure.""" + from google.adk.agents.llm.task._finish_task_tool import FINISH_TASK_TOOL_NAME + + agent = RemoteA2aAgent( + name="test_agent", + agent_card=self.agent_card, + mode="task", + ) + self.mock_context.agent_states = {} + self.mock_context.end_of_agents = {} + self.mock_context.isolation_scope = "task-1" + self.mock_session.events = [_make_dummy_task_trigger_event()] + + def set_agent_state_side_effect(agent_name, **kwargs): + if kwargs.get("end_of_agent"): + self.mock_context.end_of_agents[agent_name] = True + else: + self.mock_context.end_of_agents.pop(agent_name, None) + + self.mock_context.set_agent_state.side_effect = set_agent_state_side_effect + + with patch.object(agent, "_ensure_resolved") as mock_ensure: + mock_ensure.side_effect = Exception("Init failed") + events = [] + async for event in agent._run_async_impl(self.mock_context): + events.append(event) + + assert len(events) == 3 + assert "Failed to initialize remote A2A agent" in events[0].error_message + assert ( + events[1].content.parts[0].function_response.name + == FINISH_TASK_TOOL_NAME + ) + assert events[2].actions.end_of_agent is True + + @pytest.mark.asyncio + async def test_run_async_impl_task_mode_releases_control_on_empty_parts(self): + """Test _run_async_impl in task mode releases control when message parts are empty.""" + from google.adk.agents.llm.task._finish_task_tool import FINISH_TASK_TOOL_NAME + + agent = RemoteA2aAgent( + name="test_agent", + agent_card=self.agent_card, + mode="task", + ) + self.mock_context.agent_states = {} + self.mock_context.end_of_agents = {} + self.mock_context.isolation_scope = "task-1" + self.mock_session.events = [_make_dummy_task_trigger_event()] + + def set_agent_state_side_effect(agent_name, **kwargs): + if kwargs.get("end_of_agent"): + self.mock_context.end_of_agents[agent_name] = True + else: + self.mock_context.end_of_agents.pop(agent_name, None) + + self.mock_context.set_agent_state.side_effect = set_agent_state_side_effect + + with patch.object(agent, "_ensure_resolved"): + with patch.object( + agent, "_create_a2a_request_for_user_function_response" + ) as mock_create_func: + mock_create_func.return_value = None + with patch.object( + agent, "_construct_message_parts_from_session" + ) as mock_construct: + mock_construct.return_value = ([], None) + + events = [] + async for event in agent._run_async_impl(self.mock_context): + events.append(event) + + assert len(events) == 3 + assert events[0].content is not None + assert ( + events[1].content.parts[0].function_response.name + == FINISH_TASK_TOOL_NAME + ) + assert events[2].actions.end_of_agent is True + + @pytest.mark.asyncio + async def test_run_async_impl_a2a_http_error_in_task_mode(self): + """Test _run_async_impl task mode hand-back when A2A send_message raises HTTP error.""" + from google.adk.agents.llm.task._finish_task_tool import FINISH_TASK_ERROR_RESULT + from google.adk.agents.llm.task._finish_task_tool import FINISH_TASK_TOOL_NAME + + agent = RemoteA2aAgent( + name="test_agent", + agent_card=self.agent_card, + genai_part_converter=self.mock_genai_part_converter, + a2a_part_converter=self.mock_a2a_part_converter, + mode="task", + ) + + HTTPErrorClass = _compat.A2A_HTTP_ERRORS[0] + if _compat.IS_A2A_V1: + error_instance = HTTPErrorClass("HTTP Error 500") + else: + error_instance = HTTPErrorClass( + status_code=500, message="Internal Server Error" + ) + + self.mock_context.agent_states = {} + self.mock_context.end_of_agents = {} + self.mock_context.isolation_scope = "task-1" + self.mock_session.events = [_make_dummy_task_trigger_event()] + + def set_agent_state_side_effect(agent_name, **kwargs): + if kwargs.get("end_of_agent"): + self.mock_context.end_of_agents[agent_name] = True + else: + self.mock_context.end_of_agents.pop(agent_name, None) + + self.mock_context.set_agent_state.side_effect = set_agent_state_side_effect + + # Mock _ensure_resolved to return mock client + mock_a2a_client = Mock() + mock_send_message = AsyncMock() + mock_send_message.__aiter__.side_effect = error_instance + mock_a2a_client.send_message.return_value = mock_send_message + mock_ensure_resolved = AsyncMock(return_value=mock_a2a_client) + + with patch.object(agent, "_ensure_resolved", mock_ensure_resolved): + with patch.object( + agent, "_create_a2a_request_for_user_function_response" + ) as mock_create_func: + mock_create_func.return_value = None + + with patch.object( + agent, "_construct_message_parts_from_session" + ) as mock_construct: + mock_a2a_part = _compat.make_text_part("test") + mock_construct.return_value = ( + [mock_a2a_part], + "context-123", + ) + + agent._a2a_client = mock_a2a_client + + with patch( + "google.adk.agents.remote_a2a_agent.build_a2a_request_log" + ) as mock_req_log: + mock_req_log.return_value = "Mock request log" + + events = [] + async for event in agent._run_async_impl(self.mock_context): + events.append(event) + + # In task mode, it should yield: + # 1. The initial error event (Event with error_message) + # 2. The finish_task error event (Event with user role and finish_task FR) + # 3. The agent state event (Event with end_of_agent=True) + assert len(events) == 3 + + assert "A2A request failed" in events[0].error_message + + # The second event should be the finish_task error event + assert events[1].content is not None + fr = events[1].content.parts[0].function_response + assert fr is not None + assert fr.name == FINISH_TASK_TOOL_NAME + assert fr.response == {"result": FINISH_TASK_ERROR_RESULT} + + # The third event should be the agent state event + assert self.mock_context.end_of_agents[agent.name] is True + @pytest.mark.asyncio async def test_run_live_impl_not_implemented(self): """Test that _run_live_impl raises NotImplementedError.""" @@ -4116,6 +4881,552 @@ class TestRemoteA2aAgentDeepcopy: ) +class TestFindFinishTaskArgsFromHistory: + """Test _find_finish_task_args_from_history helper function.""" + + def test_find_finish_task_args_no_filtering(self): + # Session with multiple events + event1 = Mock(spec=Event) + event1.isolation_scope = "task-1" + event1.get_function_calls.return_value = [ + genai_types.FunctionCall( + id="fc-1", name="finish_task", args={"result": "task-1-done"} + ) + ] + + event2 = Mock(spec=Event) + event2.isolation_scope = "task-2" + event2.get_function_calls.return_value = [ + genai_types.FunctionCall( + id="fc-2", name="finish_task", args={"result": "task-2-done"} + ) + ] + + session = Mock(spec=Session) + session.events = [event1, event2] + + # Without isolation_scope, it should return the latest (event2) + args = remote_a2a_agent._find_finish_task_args_from_history(session) + assert args == {"result": "task-2-done"} + + def test_find_finish_task_args_with_filtering(self): + # Session with multiple events + event1 = Mock(spec=Event) + event1.isolation_scope = "task-1" + event1.get_function_calls.return_value = [ + genai_types.FunctionCall( + id="fc-1", name="finish_task", args={"result": "task-1-done"} + ) + ] + + event2 = Mock(spec=Event) + event2.isolation_scope = "task-2" + event2.get_function_calls.return_value = [ + genai_types.FunctionCall( + id="fc-2", name="finish_task", args={"result": "task-2-done"} + ) + ] + + session = Mock(spec=Session) + session.events = [event1, event2] + + # With isolation_scope="task-1", it should return event1's args + args = remote_a2a_agent._find_finish_task_args_from_history( + session, "task-1" + ) + assert args == {"result": "task-1-done"} + + # With isolation_scope="task-3", it should return None + args = remote_a2a_agent._find_finish_task_args_from_history( + session, "task-3" + ) + assert args is None + + def test_find_finish_task_args_with_matching_fr_id(self): + # Session with multiple finish_task FCs in the same scope + event1 = Mock(spec=Event) + event1.isolation_scope = "task-1" + event1.get_function_calls.return_value = [ + genai_types.FunctionCall( + id="fc-1", name="finish_task", args={"result": "first-attempt"} + ) + ] + + event2 = Mock(spec=Event) + event2.isolation_scope = "task-1" + event2.get_function_calls.return_value = [ + genai_types.FunctionCall( + id="fc-2", name="finish_task", args={"result": "second-attempt"} + ) + ] + + session = Mock(spec=Session) + session.events = [event1, event2] + + # Create a FR event with matching ID "fc-1" (the older one) + fr_event = Mock(spec=Event) + fr_event.get_function_responses.return_value = [ + genai_types.FunctionResponse( + id="fc-1", name="finish_task", response={"result": "SUCCESS"} + ) + ] + + # Should return event1's args because it matches fc-1, even though event2 is + # newer + args = remote_a2a_agent._find_finish_task_args_from_history( + session, "task-1", completed_fr_event=fr_event + ) + assert args == {"result": "first-attempt"} + + def test_find_finish_task_args_with_non_matching_fr_id(self): + # Session with a finish_task FC + event1 = Mock(spec=Event) + event1.isolation_scope = "task-1" + event1.get_function_calls.return_value = [ + genai_types.FunctionCall( + id="fc-1", name="finish_task", args={"result": "done"} + ) + ] + + session = Mock(spec=Session) + session.events = [event1] + + # Create a FR event with a non-matching ID "fc-different" + fr_event = Mock(spec=Event) + fr_event.get_function_responses.return_value = [ + genai_types.FunctionResponse( + id="fc-different", + name="finish_task", + response={"result": "SUCCESS"}, + ) + ] + + # Should return None because ID doesn't match + args = remote_a2a_agent._find_finish_task_args_from_history( + session, "task-1", completed_fr_event=fr_event + ) + assert args is None + + +class _TestSingleFieldOutput(BaseModel): + result: str + + +class TestRemoteA2aAgentTaskModeOutputUnwrapping: + """Test that RemoteA2aAgent correctly unwraps task output based on schema.""" + + @pytest.mark.parametrize( + "output_schema, args, expected_output", + [ + # Case 1: output_schema is None (default). Should NOT unwrap. + (None, {"result": "hello"}, {"result": "hello"}), + # Case 2: output_schema is primitive (str). Should unwrap. + (str, {"result": "hello"}, "hello"), + # Case 3: output_schema is BaseModel with 'result' field. Should NOT unwrap. + (_TestSingleFieldOutput, {"result": "hello"}, {"result": "hello"}), + # Case 4: output_schema is primitive (int). Should unwrap. + (int, {"result": 42}, 42), + # Case 5: Custom schema with multiple fields. Should NOT unwrap. + ( + dict, + {"result": "hello", "other": "world"}, + {"result": "hello", "other": "world"}, + ), + ], + ids=[ + "default_schema_no_unwrap", + "primitive_str_unwrap", + "basemodel_single_field_no_unwrap", + "primitive_int_unwrap", + "dict_no_unwrap", + ], + ) + @pytest.mark.asyncio + async def test_output_unwrapping(self, output_schema, args, expected_output): + agent_card = create_test_agent_card() + agent = RemoteA2aAgent( + name="test_agent", + agent_card=agent_card, + mode="task", + output_schema=output_schema, + ) + + mock_context = Mock(spec=InvocationContext) + mock_context.session = Mock(spec=Session) + mock_context.session.events = [_make_dummy_task_trigger_event()] + mock_context.session.state = {} + mock_context.agent_states = {} + mock_context.end_of_agents = {} + mock_context.isolation_scope = "task-1" + mock_context.invocation_id = "invocation-123" + mock_context.branch = "main" + + # Mock a2a client as regular Mock + mock_a2a_client = Mock() + mock_send_message = AsyncMock() + mock_ensure_resolved = AsyncMock(return_value=mock_a2a_client) + + # Mock _ensure_resolved to return our mock client + with patch.object(agent, "_ensure_resolved", mock_ensure_resolved): + # Mock _construct_message_parts_from_session to avoid early exit + with patch.object( + agent, "_construct_message_parts_from_session" + ) as mock_construct: + mock_a2a_part = _compat.make_text_part("test_message") + mock_construct.return_value = ([mock_a2a_part], "context-123") + + # Use a real A2AMessage wrapped in StreamResponse for the mock stream + mock_a2a_message = A2AMessage( + message_id="m1", + role=_compat.ROLE_USER, + parts=[mock_a2a_part], + context_id="context-123", + ) + mock_response = _make_stream_message(mock_a2a_message) + mock_send_message.__aiter__.return_value = [mock_response] + mock_a2a_client.send_message.return_value = mock_send_message + agent._a2a_client = mock_a2a_client + + # Mock _handle_a2a_response to return a success finish_task FR event + mock_event = Mock(spec=Event) + mock_event.custom_metadata = {} + mock_fr = genai_types.FunctionResponse( + id="ft-1", + name="finish_task", + response={"result": "Task completed."}, + ) + mock_event.get_function_responses.return_value = [mock_fr] + mock_event.get_function_calls.return_value = [] + + with patch.object( + agent, "_handle_a2a_response", new_callable=AsyncMock + ) as mock_handle: + mock_handle.return_value = mock_event + + # Mock _find_finish_task_args_from_history to return our test args + with patch( + "google.adk.agents.remote_a2a_agent._find_finish_task_args_from_history" + ) as mock_find_args: + mock_find_args.return_value = args + + events = [] + async for ev in agent._run_async_impl(mock_context): + events.append(ev) + + # We expect at least the success event + assert len(events) >= 1 + # The first yielded event should be the one modified with output + success_event = events[0] + assert success_event.output == expected_output + + @pytest.mark.asyncio + async def test_output_unwrapping_integration_with_history(self): + """Test that output unwrapping works when driving real events through the stream. + + This verifies that the finish_task FunctionCall event is correctly + placed in history before the FunctionResponse event is processed, + allowing _find_finish_task_args_from_history to find it. + """ + output_schema = str + expected_output = "hello" + agent_card = create_test_agent_card() + agent = RemoteA2aAgent( + name="test_agent", + agent_card=agent_card, + mode="task", + output_schema=output_schema, + ) + + mock_context = Mock(spec=InvocationContext) + mock_context.session = Mock(spec=Session) + mock_context.session.events = [_make_dummy_task_trigger_event()] + mock_context.session.state = {} + mock_context.agent_states = {} + mock_context.end_of_agents = {} + mock_context.isolation_scope = "task-1" + mock_context.invocation_id = "invocation-123" + mock_context.branch = "main" + + # Mock a2a client + mock_a2a_client = Mock() + mock_send_message = AsyncMock() + mock_ensure_resolved = AsyncMock(return_value=mock_a2a_client) + + # Prepare real A2A messages for FC and FR + # 1. FunctionCall for finish_task + fc_data = { + "name": "finish_task", + "args": {"result": expected_output}, + "id": "ft-1", + } + fc_meta = { + "adk_type": "function_call", + } + fc_part = _compat.make_data_part(data=fc_data, metadata=fc_meta) + fc_message = A2AMessage( + message_id="m-fc", + role=_compat.ROLE_AGENT, + parts=[fc_part], + context_id="context-123", + ) + + # 2. FunctionResponse for finish_task + fr_data = { + "name": "finish_task", + "response": {"result": FINISH_TASK_SUCCESS_RESULT}, + "id": "ft-1", + } + fr_meta = { + "adk_type": "function_response", + } + fr_part = _compat.make_data_part(data=fr_data, metadata=fr_meta) + fr_message = A2AMessage( + message_id="m-fr", + role=_compat.ROLE_USER, + parts=[fr_part], + context_id="context-123", + ) + + # Mock the stream to yield FC then FR + stream_fc = _make_stream_message(fc_message) + stream_fr = _make_stream_message(fr_message) + mock_send_message.__aiter__.return_value = [stream_fc, stream_fr] + mock_a2a_client.send_message.return_value = mock_send_message + agent._a2a_client = mock_a2a_client + + with patch.object(agent, "_ensure_resolved", mock_ensure_resolved): + # We do NOT mock _handle_a2a_response or _find_finish_task_args_from + # history + with patch.object( + agent, "_construct_message_parts_from_session" + ) as mock_construct: + mock_dummy_part = _compat.make_text_part("dummy_input") + mock_construct.return_value = ([mock_dummy_part], "context-123") + + events = [] + async for ev in agent._run_async_impl(mock_context): + events.append(ev) + # Simulate the runner setting isolation_scope and appending to history + if ev.isolation_scope is None: + ev.isolation_scope = mock_context.isolation_scope + mock_context.session.events.append(ev) + + # We expect at least the FC event and the FR event + assert len(events) >= 2 + + # Find the FR event (it should have output populated) + fr_event = None + for ev in events: + if ev.get_function_responses(): + fr_event = ev + break + + assert fr_event is not None + assert fr_event.output == expected_output + + +class TestRemoteA2aAgentTaskModeFailurePropagation: + """Test that RemoteA2aAgent propagates task failures in task mode.""" + + @pytest.mark.asyncio + async def test_fails_on_task_state_failed(self): + agent_card = create_test_agent_card() + agent = RemoteA2aAgent( + name="test_agent", + agent_card=agent_card, + mode="task", + ) + + mock_context = Mock(spec=InvocationContext) + mock_context.session = Mock(spec=Session) + mock_context.session.events = [_make_dummy_task_trigger_event()] + mock_context.session.state = {} + mock_context.agent_states = {} + mock_context.end_of_agents = {} + mock_context.isolation_scope = "task-1" + mock_context.invocation_id = "invocation-123" + mock_context.branch = "main" + + def set_agent_state_side_effect(agent_name, **kwargs): + if kwargs.get("end_of_agent"): + mock_context.end_of_agents[agent_name] = True + else: + mock_context.end_of_agents.pop(agent_name, None) + + mock_context.set_agent_state.side_effect = set_agent_state_side_effect + + # Mock a2a client as regular Mock + mock_a2a_client = Mock() + mock_send_message = AsyncMock() + mock_ensure_resolved = AsyncMock(return_value=mock_a2a_client) + + with patch.object(agent, "_ensure_resolved", mock_ensure_resolved): + with patch.object( + agent, "_construct_message_parts_from_session" + ) as mock_construct: + mock_a2a_part = _compat.make_text_part("test_message") + mock_construct.return_value = ([mock_a2a_part], "context-123") + + error_message = A2AMessage( + message_id="err-msg-1", + role=_compat.ROLE_AGENT, + parts=[_compat.make_text_part("Simulated remote task failure")], + context_id="context-123", + ) + task_status = A2ATaskStatus( + state=_compat.TS_FAILED, + message=error_message, + ) + failed_task = A2ATask( + id="task-1", + context_id="context-123", + status=task_status, + ) + + mock_response = _make_stream_task(failed_task) + mock_send_message.__aiter__.return_value = [mock_response] + mock_a2a_client.send_message.return_value = mock_send_message + agent._a2a_client = mock_a2a_client + + mock_event = Event( + author=agent.name, + invocation_id=mock_context.invocation_id, + branch=mock_context.branch, + content=genai_types.Content( + role="model", + parts=[ + genai_types.Part.from_text( + text="Simulated remote task failure" + ) + ], + ), + ) + + with patch.object( + agent, "_handle_a2a_response", new_callable=AsyncMock + ) as mock_handle: + mock_handle.return_value = mock_event + + events = [] + async for ev in agent._run_async_impl(mock_context): + events.append(ev) + + assert len(events) == 4 + assert events[0] == mock_event + assert ( + events[1].error_message + == "Remote A2A task failed: Simulated remote task failure" + ) + # Verify finish_task event + assert ( + events[2].content.parts[0].function_response.name + == FINISH_TASK_TOOL_NAME + ) + assert events[2].content.parts[0].function_response.response == { + "result": FINISH_TASK_ERROR_RESULT + } + assert events[3].actions.end_of_agent is True + + mock_context.set_agent_state.assert_called_once_with( + agent.name, end_of_agent=True + ) + + @pytest.mark.asyncio + async def test_completes_on_task_state_canceled(self): + agent_card = create_test_agent_card() + agent = RemoteA2aAgent( + name="test_agent", + agent_card=agent_card, + mode="task", + ) + + mock_context = Mock(spec=InvocationContext) + mock_context.session = Mock(spec=Session) + mock_context.session.events = [_make_dummy_task_trigger_event()] + mock_context.session.state = {} + mock_context.agent_states = {} + mock_context.end_of_agents = {} + mock_context.isolation_scope = "task-1" + mock_context.invocation_id = "invocation-123" + mock_context.branch = "main" + + def set_agent_state_side_effect(agent_name, **kwargs): + if kwargs.get("end_of_agent"): + mock_context.end_of_agents[agent_name] = True + else: + mock_context.end_of_agents.pop(agent_name, None) + + mock_context.set_agent_state.side_effect = set_agent_state_side_effect + + mock_a2a_client = Mock() + mock_send_message = AsyncMock() + mock_ensure_resolved = AsyncMock(return_value=mock_a2a_client) + + with patch.object(agent, "_ensure_resolved", mock_ensure_resolved): + with patch.object( + agent, "_construct_message_parts_from_session" + ) as mock_construct: + mock_a2a_part = _compat.make_text_part("test_message") + mock_construct.return_value = ([mock_a2a_part], "context-123") + + task_status = A2ATaskStatus( + state=_compat.TS_CANCELED, + message=None, + ) + canceled_task = A2ATask( + id="task-1", + context_id="context-123", + status=task_status, + ) + + mock_response = _make_stream_task(canceled_task) + mock_send_message.__aiter__.return_value = [mock_response] + mock_a2a_client.send_message.return_value = mock_send_message + agent._a2a_client = mock_a2a_client + + mock_event = Event( + author=agent.name, + invocation_id=mock_context.invocation_id, + branch=mock_context.branch, + content=genai_types.Content( + role="model", + parts=[genai_types.Part.from_text(text="Some progress")], + ), + ) + + with patch.object( + agent, "_handle_a2a_response", new_callable=AsyncMock + ) as mock_handle: + mock_handle.return_value = mock_event + + events = [] + async for ev in agent._run_async_impl(mock_context): + events.append(ev) + + assert len(events) == 4 + assert events[0] == mock_event + assert ( + events[1].error_message == "Remote A2A task failed: Task canceled" + ) + # Verify finish_task event + assert ( + events[2].content.parts[0].function_response.name + == FINISH_TASK_TOOL_NAME + ) + assert events[2].content.parts[0].function_response.response == { + "result": FINISH_TASK_ERROR_RESULT + } + assert events[2].output is None + assert ( + events[2].error_message == "Remote A2A task failed: Task canceled" + ) + assert events[3].actions.end_of_agent is True + assert mock_context.end_of_agents[agent.name] is True + mock_context.set_agent_state.assert_called_once_with( + agent.name, end_of_agent=True + ) + + class TestRemoteA2aAgentWorkflowOutput: """Tests that RemoteA2aAgent surfaces a workflow-node output value. @@ -4444,6 +5755,79 @@ class TestRemoteA2aAgentWorkflowOutput: assert join_outputs, "JoinNode should emit an aggregated output event" assert join_outputs[0].output == {"remote_agent": "agent reply"} + @pytest.mark.asyncio + async def test_run_impl_task_mode_failure_does_not_raise_output_already_set( + self, + ): + """Guards against ``ValueError: Output already set`` during task failure. + + In task mode, the output must not be dynamically promoted from the + original message chunk if it was already/will be set by the task + termination finish_task event. + """ + server_error_event = self._make_text_event( + text="Simulated error", task_state="failed" + ) + error_event = Event( + author="remote_agent", + error_message="Remote A2A task failed: Simulated error", + ) + finish_event = Event( + author="remote_agent", + content=genai_types.Content( + role="user", + parts=[ + genai_types.Part( + function_response=genai_types.FunctionResponse( + name=FINISH_TASK_TOOL_NAME, + response={"result": FINISH_TASK_ERROR_RESULT}, + ) + ) + ], + ), + output="Simulated error", + ) + + class _StubRemoteAgent(RemoteA2aAgent): + + async def _run_async_impl(self, ctx): + yield server_error_event + yield error_event + yield finish_event + + agent = _StubRemoteAgent( + name="remote_agent", + agent_card=create_test_agent_card(), + mode="task", + ) + agent.parent_agent = Mock() + + from google.adk.apps.app import App + from google.adk.workflow._join_node import JoinNode + from google.adk.workflow._workflow import Workflow + + from tests.unittests import testing_utils + + workflow = Workflow( + name="wf", + edges=[("START", agent, JoinNode(name="join"))], + ) + app_instance = App(name="t", root_agent=workflow) + runner = testing_utils.InMemoryRunner(app=app_instance) + + # Should run successfully without raising "Output already set" + events = await runner.run_async(testing_utils.get_user_content("start")) + + join_outputs = [ + e + for e in events + if isinstance(e, Event) + and e.output is not None + and "join" in (e.node_info.path or "") + ] + assert join_outputs + assert join_outputs[0].output == {"remote_agent": "Simulated error"} + # --------------------------------------------------------------------------- # Regression coverage for the A2A human-input resume rewrite (b/540026826) and diff --git a/tests/unittests/test_runners.py b/tests/unittests/test_runners.py index 1521999e..3d2ed1eb 100644 --- a/tests/unittests/test_runners.py +++ b/tests/unittests/test_runners.py @@ -23,12 +23,17 @@ from typing import AsyncGenerator from typing import Optional from unittest.mock import AsyncMock from unittest.mock import create_autospec +from unittest.mock import patch from google.adk import runners from google.adk.agents.base_agent import BaseAgent from google.adk.agents.context_cache_config import ContextCacheConfig from google.adk.agents.invocation_context import InvocationContext +from google.adk.agents.llm.task._finish_task_tool import FINISH_TASK_ERROR_RESULT +from google.adk.agents.llm.task._finish_task_tool import FINISH_TASK_SUCCESS_RESULT +from google.adk.agents.llm.task._finish_task_tool import FINISH_TASK_TOOL_NAME from google.adk.agents.llm_agent import LlmAgent +from google.adk.agents.remote_a2a_agent import RemoteA2aAgent from google.adk.agents.run_config import RunConfig from google.adk.apps.app import App from google.adk.apps.app import ResumabilityConfig @@ -2515,5 +2520,143 @@ def test_runner_agent_is_a_class_attribute(): assert create_autospec(Runner).agent is not None +@pytest.mark.asyncio +async def test_runner_delegation_finds_active_task_scope_on_non_terminal_error(): + """find_active_task_scope ignores non-terminal errors; remains on active task.""" + session_service = InMemorySessionService() + agent = LlmAgent(name="task_agent", mode="task") + runner = Runner( + app_name=TEST_APP_ID, agent=agent, session_service=session_service + ) + await session_service.create_session( + app_name=TEST_APP_ID, user_id=TEST_USER_ID, session_id=TEST_SESSION_ID + ) + + # Simulate non-terminal error (validation failure) + events = [ + Event( + author="task_agent", + invocation_id="inv-1", + isolation_scope="scope-1", + content=types.Content( + parts=[ + types.Part.from_function_response( + name=FINISH_TASK_TOOL_NAME, + response={"result": "Validation failed; retry"}, + ) + ] + ), + ) + ] + session = await session_service.get_session( + app_name=TEST_APP_ID, user_id=TEST_USER_ID, session_id=TEST_SESSION_ID + ) + session.events.extend(events) + assert runners._find_active_task_scope(session) == ("scope-1", "inv-1") + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "result", [FINISH_TASK_SUCCESS_RESULT, FINISH_TASK_ERROR_RESULT] +) +async def test_runner_delegation_closes_active_task_scope_on_terminal_results( + result, +): + """find_active_task_scope returns None if the task scope has finished with a terminal result.""" + session_service = InMemorySessionService() + agent = LlmAgent(name="task_agent", mode="task") + runner = Runner( + app_name=TEST_APP_ID, agent=agent, session_service=session_service + ) + await session_service.create_session( + app_name=TEST_APP_ID, user_id=TEST_USER_ID, session_id=TEST_SESSION_ID + ) + + # Simulate terminal result + events = [ + Event( + author="task_agent", + invocation_id="inv-1", + isolation_scope="scope-1", + content=types.Content( + parts=[ + types.Part.from_function_response( + name=FINISH_TASK_TOOL_NAME, + response={"result": result}, + ) + ] + ), + ) + ] + session = await session_service.get_session( + app_name=TEST_APP_ID, user_id=TEST_USER_ID, session_id=TEST_SESSION_ID + ) + session.events.extend(events) + assert runners._find_active_task_scope(session) is None + + +@pytest.mark.asyncio +async def test_runner_picks_coordinator_when_has_remote_a2a_task_subagent(): + """Runner runs coordinator, not sub-agent, when RemoteA2aAgent is in task mode.""" + session_service = InMemorySessionService() + + sub_agent = RemoteA2aAgent( + name="remote_task_agent", + agent_card="https://example.com/rpc", + mode="task", + ) + + # Coordinator LlmAgent in chat mode + coordinator = LlmAgent( + name="coordinator", mode="chat", sub_agents=[sub_agent] + ) + sub_agent.parent_agent = coordinator + + runner = Runner( + app_name=TEST_APP_ID, agent=coordinator, session_service=session_service + ) + await session_service.create_session( + app_name=TEST_APP_ID, user_id=TEST_USER_ID, session_id=TEST_SESSION_ID + ) + + # Simulate some events so _find_agent_to_run would be called on resume. + events = [ + Event( + author="remote_task_agent", + invocation_id="inv-1", + isolation_scope="scope-1", + content=types.Content(parts=[types.Part(text="task progress")]), + ) + ] + session = await session_service.get_session( + app_name=TEST_APP_ID, user_id=TEST_USER_ID, session_id=TEST_SESSION_ID + ) + session.events.extend(events) + + # Mock _run_node_async to just yield a dummy event and return + async def mock_run_node_async(*args, **kwargs): + yield Event( + author="system", + content=types.Content(parts=[types.Part(text="dummy")]), + ) + + with patch.object( + runner, "_run_node_async", side_effect=mock_run_node_async + ) as mock_run_node: + # Run with new message (resume-like) + async for _ in runner.run_async( + user_id=TEST_USER_ID, + session_id=TEST_SESSION_ID, + new_message=types.Content(parts=[types.Part(text="user reply")]), + ): + pass + + # Verify that _run_node_async was called with the coordinator (self.agent) + # not the sub_agent. + assert mock_run_node.call_count == 1 + called_node = mock_run_node.call_args[1].get("node") + assert called_node == coordinator + + if __name__ == "__main__": pytest.main([__file__]) diff --git a/tests/unittests/workflow/utils/test_workflow_graph_utils.py b/tests/unittests/workflow/utils/test_workflow_graph_utils.py index 57223b14..fcd8df3e 100644 --- a/tests/unittests/workflow/utils/test_workflow_graph_utils.py +++ b/tests/unittests/workflow/utils/test_workflow_graph_utils.py @@ -14,7 +14,10 @@ from __future__ import annotations +from unittest.mock import Mock + from google.adk.agents.llm_agent import LlmAgent +from google.adk.agents.remote_a2a_agent import RemoteA2aAgent from google.adk.tools.base_tool import BaseTool from google.adk.workflow._base_node import BaseNode from google.adk.workflow._base_node import START @@ -134,3 +137,53 @@ class TestBuildNode: standalone = LlmAgent(name="standalone", instruction="test") built_standalone = build_node(standalone) assert built_standalone.mode == "single_turn" + + def test_build_node_remote_a2a_agent_non_task(self): + """build_node does not wrap RemoteA2aAgent in task wrapper if mode is not task.""" + + class DummyRemoteAgent(RemoteA2aAgent): + + def __init__(self, mode=None): + super().__init__(name="dummy", agent_card="dummy_card", mode=mode) + self.parent_agent = None + + def clone(self, *args, **kwargs): + raise AssertionError("clone should not be called") + + agent = DummyRemoteAgent(mode=None) + built = build_node(agent) + assert built == agent + + def test_build_node_remote_a2a_agent_task(self): + """build_node raises ValueError if RemoteA2aAgent mode is task.""" + + class DummyRemoteAgent(RemoteA2aAgent): + + def __init__(self, mode="task"): + super().__init__(name="dummy", agent_card="dummy_card", mode=mode) + self.parent_agent = None + + agent = DummyRemoteAgent(mode="task") + with pytest.raises( + ValueError, + match=( + "RemoteA2aAgent in task mode is not supported as a standalone" + " workflow node. It is only supported in tool-delegation mode." + ), + ): + build_node(agent) + + def test_build_node_remote_a2a_agent_task_with_parent(self): + """build_node allows task-mode RemoteA2aAgent if parent_agent is set.""" + + class DummyRemoteAgent(RemoteA2aAgent): + + def __init__(self, mode="task"): + super().__init__(name="dummy", agent_card="dummy_card", mode=mode) + self.parent_agent = Mock() + + agent = DummyRemoteAgent(mode="task") + built = build_node(agent) + assert built is not agent + assert built.mode == "task" + assert built.wait_for_output is True