Files
google--adk-python/tests/unittests/a2a/integration/test_client_server.py
T
wehub-resource-sync ec2b666284
Continuous Integration / Pre-commit Linter (push) Has been cancelled
Continuous Integration / Mypy Check (Python 3.10) (push) Has been cancelled
Continuous Integration / Mypy Check (Python 3.11) (push) Has been cancelled
Continuous Integration / Mypy Check (Python 3.12) (push) Has been cancelled
Continuous Integration / Mypy Check (Python 3.13) (push) Has been cancelled
Continuous Integration / Unit Tests (Python 3.10) (push) Has been cancelled
Continuous Integration / Unit Tests (Python 3.11) (push) Has been cancelled
Continuous Integration / Unit Tests (Python 3.12) (push) Has been cancelled
Continuous Integration / Unit Tests (Python 3.13) (push) Has been cancelled
Continuous Integration / Unit Tests (Python 3.14) (push) Has been cancelled
Continuous Integration / A2A v0.3 Tests (Python 3.10) (push) Has been cancelled
Continuous Integration / A2A v0.3 Tests (Python 3.11) (push) Has been cancelled
Continuous Integration / A2A v0.3 Tests (Python 3.12) (push) Has been cancelled
Copybara PR Handler / close-imported-pr (push) Has been cancelled
Continuous Integration / A2A v0.3 Tests (Python 3.13) (push) Has been cancelled
Continuous Integration / A2A v0.3 Tests (Python 3.14) (push) Has been cancelled
chore: import upstream snapshot with attribution
2026-07-13 13:25:13 +08:00

822 lines
25 KiB
Python

# Copyright 2026 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Integration tests for A2A client-server interaction."""
import logging
from unittest.mock import AsyncMock
try:
from a2a.server.apps.jsonrpc.fastapi_app import A2AFastAPIApplication
from a2a.server.request_handlers.request_handler import RequestHandler
except ImportError:
A2AFastAPIApplication = None
RequestHandler = None
from a2a.types import Message as A2AMessage
from a2a.types import Part as A2APart
from a2a.types import Task
from a2a.types import TaskStatus
try:
# 0.3.x-only wrapper part type; these integration tests are skipped on 1.x.
from a2a.types import TextPart
except ImportError:
TextPart = None
from google.adk.a2a import _compat
from google.adk.a2a.agent.interceptors.new_integration_extension import _NEW_A2A_ADK_INTEGRATION_EXTENSION
from google.adk.a2a.converters.to_adk_event import MOCK_FUNCTION_CALL_FOR_REQUIRED_USER_INPUT
from google.adk.a2a.executor.config import A2aAgentExecutorConfig
from google.adk.a2a.executor.interceptors.include_artifacts_in_a2a_event import include_artifacts_in_a2a_event_interceptor
from google.adk.agents.remote_a2a_agent import A2A_METADATA_PREFIX
from google.adk.events.event import Event
from google.adk.events.event_actions import EventActions
from google.adk.platform import uuid as platform_uuid
from google.adk.runners import Runner
from google.adk.sessions.in_memory_session_service import InMemorySessionService
from google.genai import types
import pytest
pytestmark = pytest.mark.skipif(
_compat.IS_A2A_V1,
reason="integration tests use 0.3-only A2AFastAPIApplication",
)
from .client import create_a2a_client
from .client import create_client
from .server import agent_card
from .server import create_server_app
logger = logging.getLogger("google_adk." + __name__)
def create_streaming_mock_run_async(received_requests: list):
"""Creates a mock_run_async that streams multiple chunks."""
async def mock_run_async(**kwargs):
received_requests.append(kwargs)
yield Event(
author="FakeAgent",
content=types.Content(parts=[types.Part(text="Hello")]),
partial=True,
)
yield Event(
author="FakeAgent",
content=types.Content(parts=[types.Part(text=" world")]),
partial=True,
)
yield Event(
author="FakeAgent",
partial=True,
actions=EventActions(artifact_delta={"file1": 1}),
)
yield Event(
author="FakeAgent",
content=types.Content(parts=[types.Part(text="Hello world")]),
partial=False,
)
return mock_run_async
def create_non_streaming_mock_run_async(received_requests: list):
"""Creates a mock_run_async that returns a single non-streaming event."""
async def mock_run_async(**kwargs):
received_requests.append(kwargs)
yield Event(
author="FakeAgent",
content=types.Content(parts=[types.Part(text="Hello world")]),
partial=False,
)
return mock_run_async
@pytest.mark.asyncio
async def test_streaming_adk_to_streaming_a2a():
"""Test streaming of normal text chunks."""
received_requests = []
mock_run_async = create_streaming_mock_run_async(received_requests)
app = create_server_app(mock_run_async)
agent = create_client(app, streaming=True)
session_service = InMemorySessionService()
await session_service.create_session(
app_name="ClientApp", user_id="test_user", session_id="test_session"
)
client_runner = Runner(
app_name="ClientApp",
agent=agent,
session_service=session_service,
)
new_message = types.Content(parts=[types.Part(text="Hi")], role="user")
texts = []
actions = []
async for event in client_runner.run_async(
user_id="test_user", session_id="test_session", new_message=new_message
):
if event.content and event.content.parts:
for p in event.content.parts:
if p.text:
texts.append(p.text)
if event.actions and event.actions.artifact_delta:
actions.append(event.actions)
assert len(received_requests) == 1
assert received_requests[0]["session_id"] is not None
assert texts == ["Hello", " world", "Hello world"]
assert len(actions) == 1
assert actions[0].artifact_delta == {"file1": 1}
@pytest.mark.asyncio
async def test_streaming_adk_to_non_streaming_a2a():
"""Test ADK streaming into A2A Non-Streaming."""
received_requests = []
mock_run_async = create_streaming_mock_run_async(received_requests)
app = create_server_app(mock_run_async)
agent = create_client(app, streaming=False)
session_service = InMemorySessionService()
await session_service.create_session(
app_name="ClientApp", user_id="test_user", session_id="test_session"
)
client_runner = Runner(
app_name="ClientApp", agent=agent, session_service=session_service
)
new_message = types.Content(parts=[types.Part(text="Hi")], role="user")
texts = []
async for event in client_runner.run_async(
user_id="test_user", session_id="test_session", new_message=new_message
):
if event.content and event.content.parts:
for p in event.content.parts:
if p.text:
texts.append(p.text)
assert len(received_requests) == 1
assert texts == ["Hello world"]
@pytest.mark.asyncio
async def test_non_streaming_adk_to_streaming_a2a():
"""Test ADK Non-Streaming into A2A Streaming."""
received_requests = []
mock_run_async = create_non_streaming_mock_run_async(received_requests)
app = create_server_app(mock_run_async)
agent = create_client(app, streaming=True)
session_service = InMemorySessionService()
await session_service.create_session(
app_name="ClientApp", user_id="test_user", session_id="test_session"
)
client_runner = Runner(
app_name="ClientApp", agent=agent, session_service=session_service
)
new_message = types.Content(parts=[types.Part(text="Hi")], role="user")
texts = []
async for event in client_runner.run_async(
user_id="test_user", session_id="test_session", new_message=new_message
):
if event.content and event.content.parts:
for p in event.content.parts:
if p.text:
texts.append(p.text)
assert len(received_requests) == 1
assert texts == ["Hello world"]
@pytest.mark.asyncio
async def test_non_streaming_adk_to_non_streaming_a2a():
"""Test ADK Non-Streaming into A2A Non-Streaming."""
received_requests = []
mock_run_async = create_non_streaming_mock_run_async(received_requests)
app = create_server_app(mock_run_async)
agent = create_client(app, streaming=False)
session_service = InMemorySessionService()
await session_service.create_session(
app_name="ClientApp", user_id="test_user", session_id="test_session"
)
client_runner = Runner(
app_name="ClientApp", agent=agent, session_service=session_service
)
new_message = types.Content(parts=[types.Part(text="Hi")], role="user")
texts = []
async for event in client_runner.run_async(
user_id="test_user", session_id="test_session", new_message=new_message
):
if event.content and event.content.parts:
for p in event.content.parts:
if p.text:
texts.append(p.text)
assert len(received_requests) == 1
assert texts == ["Hello world"]
def create_streaming_mock_run_async_with_multiple_agents(
received_requests: list,
):
"""Creates a mock_run_async that streams multiple chunks."""
async def mock_run_async(**kwargs):
received_requests.append(kwargs)
yield Event(
author="FakeAgent1",
content=types.Content(parts=[types.Part(text="Hello")]),
partial=True,
)
yield Event(
author="FakeAgent2",
content=types.Content(parts=[types.Part(text=" Hi")]),
partial=True,
)
yield Event(
author="FakeAgent1",
content=types.Content(parts=[types.Part(text=" world")]),
partial=True,
)
yield Event(
author="FakeAgent2",
content=types.Content(parts=[types.Part(text=" human")]),
partial=True,
)
yield Event(
author="FakeAgent1",
content=types.Content(parts=[types.Part(text="Hello world")]),
partial=False,
)
yield Event(
author="FakeAgent2",
content=types.Content(parts=[types.Part(text="Hi human")]),
partial=False,
)
return mock_run_async
@pytest.mark.asyncio
async def test_multiple_agents_streaming_adk_to_streaming_a2a():
"""Test streaming multiple agents chunks into A2A Streaming."""
received_requests = []
mock_run_async = create_streaming_mock_run_async_with_multiple_agents(
received_requests
)
app = create_server_app(mock_run_async)
agent = create_client(app, streaming=True)
session_service = InMemorySessionService()
await session_service.create_session(
app_name="ClientApp", user_id="test_user", session_id="test_session"
)
client_runner = Runner(
app_name="ClientApp", agent=agent, session_service=session_service
)
new_message = types.Content(parts=[types.Part(text="Hi")], role="user")
texts = []
async for event in client_runner.run_async(
user_id="test_user", session_id="test_session", new_message=new_message
):
if event.content and event.content.parts:
for p in event.content.parts:
if p.text:
texts.append(p.text)
assert len(received_requests) == 1
assert texts == [
"Hello",
" Hi",
" world",
" human",
"Hello world",
"Hi human",
]
@pytest.mark.asyncio
async def test_function_calls():
"""Test function call execution from agent."""
received_requests = []
async def mock_run_async(**kwargs):
received_requests.append(kwargs)
yield Event(
author="FakeAgent",
content=types.Content(
parts=[
types.Part(
function_call=types.FunctionCall(
name="get_weather",
args={"location": "San Francisco"},
id="call_1",
)
),
types.Part(
function_response=types.FunctionResponse(
name="get_weather",
response={"temperature": "22C"},
id="call_1",
)
),
],
role="model",
),
)
app = create_server_app(mock_run_async)
agent = create_client(app)
session_service = InMemorySessionService()
await session_service.create_session(
app_name="ClientApp", user_id="test_user", session_id="test_session"
)
client_runner = Runner(
app_name="ClientApp",
agent=agent,
session_service=session_service,
)
new_message = types.Content(parts=[types.Part(text="Hi")], role="user")
func_calls = []
func_responses = []
async for event in client_runner.run_async(
user_id="test_user", session_id="test_session", new_message=new_message
):
func_calls.extend(event.get_function_calls())
if event.content and event.content.parts:
for p in event.content.parts:
if p.function_response:
func_responses.append(p.function_response)
assert len(func_calls) == 1
assert func_calls[0].name == "get_weather"
assert func_calls[0].args == {"location": "San Francisco"}
assert len(func_responses) == 1
assert func_responses[0].name == "get_weather"
assert func_responses[0].response == {"temperature": "22C"}
def create_long_running_mock_run_async(received_requests: list):
"""Creates a mock_run_async for long running function tests."""
async def mock_run_async(**kwargs):
received_requests.append(kwargs)
if len(received_requests) == 1:
yield Event(
author="FakeAgent",
content=types.Content(
parts=[
types.Part(
function_call=types.FunctionCall(
name="long_task", args={}, id="call_long"
)
)
],
role="model",
),
long_running_tool_ids={"call_long"},
)
yield Event(
author="FakeAgent",
content=types.Content(
parts=[
types.Part(
function_response=types.FunctionResponse(
name="long_task",
response={"status": "pending"},
id="call_long",
)
)
],
role="model",
),
)
else:
yield Event(
author="FakeAgent",
content=types.Content(
parts=[types.Part(text="Task completed well")], role="model"
),
)
return mock_run_async
@pytest.mark.asyncio
async def test_long_running_function_calls_success():
"""Test long running function calls flow success with user response."""
received_requests = []
mock_run_async = create_long_running_mock_run_async(received_requests)
app = create_server_app(mock_run_async)
agent = create_client(app, streaming=True)
session_service = InMemorySessionService()
await session_service.create_session(
app_name="ClientApp", user_id="test_user", session_id="test_session"
)
client_runner = Runner(
app_name="ClientApp",
agent=agent,
session_service=session_service,
)
new_message_1 = types.Content(parts=[types.Part(text="Hi")], role="user")
func_calls_1 = []
func_responses_1 = []
task_id_1 = ""
has_long_running_id = False
async for event in client_runner.run_async(
user_id="test_user", session_id="test_session", new_message=new_message_1
):
if event.custom_metadata:
task_id_1 = event.custom_metadata.get(
A2A_METADATA_PREFIX + "task_id", task_id_1
)
if (
event.long_running_tool_ids
and "call_long" in event.long_running_tool_ids
):
has_long_running_id = True
func_calls_1.extend(event.get_function_calls())
if event.content and event.content.parts:
for p in event.content.parts:
if p.function_response:
func_responses_1.append(p.function_response)
assert has_long_running_id
assert len(func_calls_1) == 1
assert func_calls_1[0].name == "long_task"
assert len(func_responses_1) == 1
assert func_responses_1[0].name == "long_task"
assert func_responses_1[0].response == {"status": "pending"}
new_message_2 = types.Content(
parts=[
types.Part(
function_response=types.FunctionResponse(
name="long_task", response={"result": "done"}, id="call_long"
)
)
],
role="user",
)
texts = []
task_id_2 = ""
async for event in client_runner.run_async(
user_id="test_user", session_id="test_session", new_message=new_message_2
):
if event.custom_metadata:
task_id_2 = event.custom_metadata.get(
A2A_METADATA_PREFIX + "task_id", task_id_2
)
if event.content and event.content.parts:
for p in event.content.parts:
if p.text:
texts.append(p.text)
assert task_id_1 == task_id_2
assert "Task completed well" in texts
@pytest.mark.asyncio
async def test_long_running_function_calls_error():
"""Test long running function calls returns error on missing response."""
received_requests = []
mock_run_async = create_long_running_mock_run_async(received_requests)
app = create_server_app(mock_run_async)
a2a_client = create_a2a_client(app, streaming=False)
request_1 = A2AMessage(
message_id=platform_uuid.new_uuid(),
parts=[A2APart(root=TextPart(text="Hi"))],
role="user",
)
response_1_events = []
async for event in a2a_client.send_message(request=request_1):
response_1_events.append(event)
assert len(response_1_events) == 1
# Extract task_id from Turn 1 responses
assert response_1_events[0][1] is None
task = response_1_events[0][0]
assert isinstance(task, Task)
assert task.status.state == _compat.TS_INPUT_REQUIRED
extracted_task_id = task.id
assert extracted_task_id is not None
request_2 = A2AMessage(
message_id=platform_uuid.new_uuid(),
parts=[A2APart(root=TextPart(text="Any update?"))],
role="user",
task_id=extracted_task_id,
context_id=task.context_id if hasattr(task, "context_id") else None,
)
response_2_events = []
async for event in a2a_client.send_message(request=request_2):
response_2_events.append(event)
# Verify that we get an error response for the second request due to missing function response
assert len(response_2_events) == 1
assert response_2_events[0][1] is None
error_response = response_2_events[0][0]
assert isinstance(error_response, Task)
assert error_response.status.message.parts[0].root.text == (
"It was not provided a function response for the function call."
)
@pytest.mark.asyncio
async def test_user_follow_up():
"""Test multi-turn interaction or follow up with state."""
received_requests = []
async def mock_run_async(**kwargs):
received_requests.append(kwargs)
# Yield response with custom metadata to test passing back
yield Event(
author="FakeAgent",
content=types.Content(
parts=[types.Part(text="Follow up response")], role="model"
),
custom_metadata={"server_state": "active"},
)
app = create_server_app(mock_run_async)
agent = create_client(app)
session_service = InMemorySessionService()
await session_service.create_session(
app_name="ClientApp", user_id="test_user", session_id="test_session"
)
client_runner = Runner(
app_name="ClientApp",
agent=agent,
session_service=session_service,
)
# First Turn
new_message_1 = types.Content(parts=[types.Part(text="Turn 1")], role="user")
async for _ in client_runner.run_async(
user_id="test_user", session_id="test_session", new_message=new_message_1
):
pass
# Second Turn
new_message_2 = types.Content(parts=[types.Part(text="Turn 2")], role="user")
last_event = None
async for event in client_runner.run_async(
user_id="test_user", session_id="test_session", new_message=new_message_2
):
last_event = event
assert len(received_requests) == 2
# The second request should carry the same session ID as the first
assert (
received_requests[1]["session_id"] == received_requests[0]["session_id"]
)
assert last_event is not None
@pytest.mark.asyncio
async def test_include_artifacts_in_a2a_event():
"""Test that artifacts are included in A2A events when the interceptor is enabled."""
async def mock_run_async(**kwargs):
yield Event(
actions=EventActions(artifact_delta={"artifact1": 1, "artifact2": 1}),
author="agent",
content=types.Content(
parts=[types.Part(text="Here are the artifacts")]
),
)
config = A2aAgentExecutorConfig(
execute_interceptors=[include_artifacts_in_a2a_event_interceptor]
)
built_app = create_server_app(mock_run_async, config=config)
a2a_client = create_a2a_client(built_app, streaming=False)
request = A2AMessage(
message_id="test_message_id",
parts=[A2APart(root=TextPart(text="Hi"))],
role="user",
)
events = []
async for event in a2a_client.send_message(request=request):
events.append(event)
assert len(events) == 1
task = events[0][0]
assert isinstance(task, Task)
assert task.artifacts is not None
assert len(task.artifacts) == 3
assert task.artifacts[0].parts[0].root.text == "Here are the artifacts"
assert task.artifacts[1].artifact_id == "artifact1_1"
assert task.artifacts[1].name == "artifact1"
assert task.artifacts[1].parts[0].root.text == "artifact content"
assert task.artifacts[2].artifact_id == "artifact2_1"
assert task.artifacts[2].name == "artifact2"
assert task.artifacts[2].parts[0].root.text == "artifact content"
@pytest.mark.asyncio
async def test_user_follow_up_sends_task_id_with_input_required():
"""Test that client follow-up sends the same task_id."""
task_id = "mocked-task-id-123"
context_id = "mocked-context-id-456"
mock_task = Task(
id=task_id,
context_id=context_id,
kind="task",
status=TaskStatus(
state=_compat.TS_INPUT_REQUIRED,
message=A2AMessage(
message_id="mocked-message-id-789",
role="user",
parts=[A2APart(root=TextPart(text="Input required"))],
),
),
metadata={_NEW_A2A_ADK_INTEGRATION_EXTENSION: True},
)
mock_handler = AsyncMock(spec=RequestHandler)
# First call returns input_required, second call completes
mock_handler.on_message_send.side_effect = [
mock_task,
Task(
id=task_id,
context_id=context_id,
kind="task",
status=TaskStatus(state=_compat.TS_COMPLETED),
metadata={_NEW_A2A_ADK_INTEGRATION_EXTENSION: True},
),
]
app = A2AFastAPIApplication(
agent_card=agent_card, http_handler=mock_handler
).build()
agent = create_client(app, streaming=False)
session_service = InMemorySessionService()
await session_service.create_session(
app_name="ClientApp", user_id="test_user", session_id="test_session"
)
client_runner = Runner(
app_name="ClientApp", agent=agent, session_service=session_service
)
# First Turn
new_message_1 = types.Content(parts=[types.Part(text="Turn 1")], role="user")
found_call_id = None
async for event in client_runner.run_async(
user_id="test_user", session_id="test_session", new_message=new_message_1
):
for call in event.get_function_calls():
if call.name == MOCK_FUNCTION_CALL_FOR_REQUIRED_USER_INPUT:
found_call_id = call.id
assert found_call_id is not None
# Second Turn (Follow-up)
function_response = types.FunctionResponse(
id=found_call_id,
name=MOCK_FUNCTION_CALL_FOR_REQUIRED_USER_INPUT,
response={"result": "Turn 2"},
)
new_message_2 = types.Content(
parts=[types.Part(function_response=function_response)], role="user"
)
async for _ in client_runner.run_async(
user_id="test_user", session_id="test_session", new_message=new_message_2
):
pass
assert mock_handler.on_message_send.call_count == 2
# Second call args
call_args_2 = mock_handler.on_message_send.call_args_list[1]
params_2 = call_args_2[0][0]
assert params_2.message.task_id == task_id
@pytest.mark.asyncio
async def test_user_follow_up_sends_task_id_with_input_required_legacy_impl():
"""Test that client follow-up sends the same task_id."""
task_id = "mocked-task-id-123"
context_id = "mocked-context-id-456"
mock_task = Task(
id=task_id,
context_id=context_id,
kind="task",
status=TaskStatus(
state=_compat.TS_INPUT_REQUIRED,
message=A2AMessage(
message_id="mocked-message-id-789",
role="user",
parts=[A2APart(root=TextPart(text="Input required"))],
),
),
)
mock_handler = AsyncMock(spec=RequestHandler)
# First call returns input_required, second call completes
mock_handler.on_message_send.side_effect = [
mock_task,
Task(
id=task_id,
context_id=context_id,
kind="task",
status=TaskStatus(state=_compat.TS_COMPLETED),
),
]
app = A2AFastAPIApplication(
agent_card=agent_card, http_handler=mock_handler
).build()
agent = create_client(app, streaming=False)
session_service = InMemorySessionService()
await session_service.create_session(
app_name="ClientApp", user_id="test_user", session_id="test_session"
)
client_runner = Runner(
app_name="ClientApp", agent=agent, session_service=session_service
)
# First Turn
new_message_1 = types.Content(parts=[types.Part(text="Turn 1")], role="user")
found_call_id = None
async for event in client_runner.run_async(
user_id="test_user", session_id="test_session", new_message=new_message_1
):
for call in event.get_function_calls():
if call.name == MOCK_FUNCTION_CALL_FOR_REQUIRED_USER_INPUT:
found_call_id = call.id
assert found_call_id is not None
# Second Turn (Follow-up)
function_response = types.FunctionResponse(
id=found_call_id,
name=MOCK_FUNCTION_CALL_FOR_REQUIRED_USER_INPUT,
response={"result": "Turn 2"},
)
new_message_2 = types.Content(
parts=[types.Part(function_response=function_response)], role="user"
)
async for _ in client_runner.run_async(
user_id="test_user", session_id="test_session", new_message=new_message_2
):
pass
assert mock_handler.on_message_send.call_count == 2
# Second call args
call_args_2 = mock_handler.on_message_send.call_args_list[1]
params_2 = call_args_2[0][0]
assert params_2.message.task_id == task_id