Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions src/google/adk/agents/base_agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -176,6 +176,11 @@ class MyAgent(BaseAgent):
will be appended to event history as an additional agent response.
"""

async def _get_transfer_description(self, ctx: InvocationContext) -> str:
"""Returns the description used to select this transfer target."""
del ctx
return self.description

def _load_agent_state(
self,
ctx: InvocationContext,
Expand Down
34 changes: 31 additions & 3 deletions src/google/adk/agents/remote_a2a_agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,7 @@
from google.adk.platform import uuid as platform_uuid
from google.genai import types as genai_types
import httpx
from typing_extensions import override

from ..a2a import _compat

Expand Down Expand Up @@ -503,6 +504,31 @@ async def _validate_agent_card(self, agent_card: AgentCard) -> None:

self._validate_card_rpc_targets(agent_card)

@override
async def _get_transfer_description(self, ctx: InvocationContext) -> str:
"""Returns local or agent-card metadata for transfer selection."""
if self.description:
return self.description

if self._agent_card:
return self._agent_card.description or ""

agent_card = await self._resolve_agent_card(ctx)
await self._validate_agent_card(agent_card)

# Public cards are shared across invocations, matching the existing client
# cache. Authenticated cards remain invocation-scoped because their metadata
# may vary by session.
per_invocation_card = bool(
self._config.card_request_interceptors
and self._agent_card_source
and self._agent_card_source.startswith(("http://", "https://"))
)
if not per_invocation_card:
self._agent_card = agent_card

return agent_card.description or ""

def _validate_card_rpc_targets(self, agent_card: AgentCard) -> None:
"""Constrains where a card fetched over the network may aim RPC traffic.

Expand Down Expand Up @@ -594,9 +620,11 @@ async def _ensure_resolved(
# Validate agent card
await self._validate_agent_card(self._agent_card)

# Update description if empty
if not self.description and self._agent_card.description:
self.description = self._agent_card.description
# A public card may already have been resolved for transfer selection.
# Preserve the existing behavior of adopting its description when the
# remote agent itself is initialized.
if not self.description and self._agent_card.description:
self.description = self._agent_card.description

# Initialize A2A client
if not self._a2a_client:
Expand Down
42 changes: 40 additions & 2 deletions src/google/adk/flows/llm_flows/agent_transfer.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,8 @@

from __future__ import annotations

import asyncio
import logging
import typing
from typing import AsyncGenerator
from typing import Sequence
Expand All @@ -35,6 +37,9 @@
from ...agents.llm_agent import LlmAgent


logger = logging.getLogger('google_adk.' + __name__)


class _AgentTransferLlmRequestProcessor(BaseLlmRequestProcessor):
"""Agent transfer request processor."""

Expand All @@ -54,11 +59,16 @@ async def run_async(
agent_names=[agent.name for agent in transfer_targets]
)

transfer_target_infos = await asyncio.gather(*[
_build_transfer_target_info(target, invocation_context)
for target in transfer_targets
])

llm_request.append_instructions([
_build_transfer_instructions(
transfer_to_agent_tool.name,
agent,
transfer_targets,
transfer_target_infos,
)
])

Expand All @@ -79,6 +89,34 @@ class _AgentLike(typing.Protocol):
description: str


class _TransferTargetInfo:
"""Invocation-scoped metadata used to build transfer instructions."""

def __init__(self, *, name: str, description: str) -> None:
self.name = name
self.description = description


async def _build_transfer_target_info(
target_agent: BaseAgent,
ctx: InvocationContext,
) -> _TransferTargetInfo:
"""Builds transfer metadata without mutating invocation-scoped values."""
try:
description = await target_agent._get_transfer_description(ctx)
except Exception as e:
logger.warning(
'Failed to load transfer description for agent %s: %s',
target_agent.name,
e,
)
description = target_agent.description
return _TransferTargetInfo(
name=target_agent.name,
description=description,
)


def _build_target_agents_info(target_agent: _AgentLike) -> str:
return f"""
Agent name: {target_agent.name}
Expand Down Expand Up @@ -134,7 +172,7 @@ def _build_transfer_instruction_body(
def _build_transfer_instructions(
tool_name: str,
agent: LlmAgent,
target_agents: Sequence[BaseAgent],
target_agents: Sequence[_AgentLike],
) -> str:
"""Build instructions for agent transfer (agent-tree variant).

Expand Down
17 changes: 17 additions & 0 deletions tests/unittests/agents/test_remote_a2a_agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -1052,6 +1052,23 @@ async def test_ensure_resolved_with_url_source(self):
assert agent._agent_card == agent_card
assert agent.description == agent_card.description

@pytest.mark.asyncio
async def test_ensure_resolved_adopts_prefetched_card_description(self):
"""Initializing a prefetched public card still adopts its description."""
agent_card = create_test_agent_card()
agent = RemoteA2aAgent(
name="test_agent", agent_card="https://example.com/agent.json"
)
agent._agent_card = agent_card
mock_factory = Mock()
mock_factory.create.return_value = Mock()
agent._a2a_client_factory = mock_factory

with patch.object(agent, "_ensure_httpx_client", new_callable=AsyncMock):
await agent._ensure_resolved(Mock())

assert agent.description == agent_card.description

@pytest.mark.asyncio
async def test_ensure_resolved_already_resolved(self):
"""Test _ensure_resolved when already resolved."""
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -20,10 +20,17 @@
"""

from typing import AsyncGenerator
from unittest.mock import AsyncMock
from unittest.mock import patch

from a2a.types import AgentCard
from google.adk.a2a import _compat
from google.adk.a2a.agent.config import A2aRemoteAgentConfig
from google.adk.a2a.agent.config import CardRequestInterceptor
from google.adk.agents.base_agent import BaseAgent
from google.adk.agents.invocation_context import InvocationContext
from google.adk.agents.llm_agent import Agent
from google.adk.agents.remote_a2a_agent import RemoteA2aAgent
from google.adk.artifacts.in_memory_artifact_service import InMemoryArtifactService
from google.adk.events.event import Event
from google.adk.flows.llm_flows import agent_transfer
Expand All @@ -47,6 +54,16 @@ async def _run_async_impl(
yield Event(author=self.name, invocation_id=ctx.invocation_id)


def _make_agent_card(description: str) -> AgentCard:
return _compat.build_agent_card(
name='remote_agent',
description=description,
version='1.0',
url='https://example.com/rpc',
protocol_binding='JSONRPC',
)


async def create_test_invocation_context(agent: Agent) -> InvocationContext:
"""Helper to create constructed InvocationContext."""
session_service = InMemorySessionService()
Expand Down Expand Up @@ -342,3 +359,150 @@ async def test_agent_transfer_with_non_llm_peer_agent():

instructions = llm_request.config.system_instruction
assert 'non_llm_peer' in instructions


@pytest.mark.asyncio
async def test_agent_transfer_loads_remote_card_description_before_selection():
"""A remote card description is available for the first transfer decision."""
mock_model = testing_utils.MockModel.create(responses=[])
remote_agent = RemoteA2aAgent(
name='remote_agent',
agent_card='https://example.com/agent-card.json',
)
main_agent = Agent(
name='main_agent',
model=mock_model,
sub_agents=[remote_agent],
)
invocation_context = await create_test_invocation_context(main_agent)
llm_request = LlmRequest()

with patch.object(
remote_agent,
'_resolve_agent_card',
new=AsyncMock(return_value=_make_agent_card('Handles remote research.')),
):
async for _ in agent_transfer.request_processor.run_async(
invocation_context, llm_request
):
pass

instructions = llm_request.config.system_instruction
assert 'Agent description: Handles remote research.' in instructions


@pytest.mark.asyncio
async def test_agent_transfer_prefers_explicit_remote_description():
"""An explicit remote description remains authoritative for transfers."""
mock_model = testing_utils.MockModel.create(responses=[])
remote_agent = RemoteA2aAgent(
name='remote_agent',
agent_card='https://example.com/agent-card.json',
description='Locally configured routing description.',
)
main_agent = Agent(
name='main_agent',
model=mock_model,
sub_agents=[remote_agent],
)
invocation_context = await create_test_invocation_context(main_agent)
llm_request = LlmRequest()

with patch.object(
remote_agent, '_resolve_agent_card', new_callable=AsyncMock
) as resolve_card:
async for _ in agent_transfer.request_processor.run_async(
invocation_context, llm_request
):
pass

instructions = llm_request.config.system_instruction
assert (
'Agent description: Locally configured routing description.'
in instructions
)
resolve_card.assert_not_awaited()


@pytest.mark.asyncio
async def test_agent_transfer_keeps_authenticated_descriptions_per_invocation():
"""Authenticated card descriptions do not leak between invocations."""
mock_model = testing_utils.MockModel.create(responses=[])
remote_agent = RemoteA2aAgent(
name='remote_agent',
agent_card='https://example.com/agent-card.json',
config=A2aRemoteAgentConfig(
card_request_interceptors=[CardRequestInterceptor()]
),
)
main_agent = Agent(
name='main_agent',
model=mock_model,
sub_agents=[remote_agent],
)
first_context = await create_test_invocation_context(main_agent)
second_context = await create_test_invocation_context(main_agent)
first_request = LlmRequest()
second_request = LlmRequest()

with patch.object(
remote_agent,
'_resolve_agent_card',
new=AsyncMock(
side_effect=[
_make_agent_card('First session description.'),
_make_agent_card('Second session description.'),
]
),
):
async for _ in agent_transfer.request_processor.run_async(
first_context, first_request
):
pass
async for _ in agent_transfer.request_processor.run_async(
second_context, second_request
):
pass

first_instructions = first_request.config.system_instruction
second_instructions = second_request.config.system_instruction
assert 'Agent description: First session description.' in first_instructions
assert 'Second session description.' not in first_instructions
assert 'Agent description: Second session description.' in second_instructions
assert 'First session description.' not in second_instructions
assert remote_agent.description == ''


@pytest.mark.asyncio
async def test_agent_transfer_continues_when_remote_card_is_unavailable(caplog):
"""An unavailable remote card does not block the parent model request."""
mock_model = testing_utils.MockModel.create(responses=[])
remote_agent = RemoteA2aAgent(
name='remote_agent',
agent_card='https://example.com/agent-card.json',
)
main_agent = Agent(
name='main_agent',
model=mock_model,
sub_agents=[remote_agent],
)
invocation_context = await create_test_invocation_context(main_agent)
llm_request = LlmRequest()

with patch.object(
remote_agent,
'_resolve_agent_card',
new=AsyncMock(side_effect=OSError('card service unavailable')),
):
async for _ in agent_transfer.request_processor.run_async(
invocation_context, llm_request
):
pass

instructions = llm_request.config.system_instruction
assert 'Agent name: remote_agent' in instructions
assert 'Agent description: ' in instructions
assert (
'Failed to load transfer description for agent remote_agent'
in caplog.text
)