Skip to content
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
16 commits
Select commit Hold shift + click to select a range
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
10 changes: 9 additions & 1 deletion src/a2a/client/base_client.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,3 @@
from collections.abc import AsyncIterator

from a2a.client.client import (
Expand Down Expand Up @@ -46,6 +46,7 @@
self,
request: Message,
*,
configuration: MessageSendConfiguration | None = None,
context: ClientCallContext | None = None,
) -> AsyncIterator[ClientEvent | Message]:
"""Sends a message to the agent.
Expand All @@ -56,12 +57,13 @@

Args:
request: The message to send to the agent.
configuration: Optional per-call overrides for message sending behavior.
context: The client call context.

Yields:
An async iterator of `ClientEvent` or a final `Message` response.
"""
config = MessageSendConfiguration(
base_config = MessageSendConfiguration(
accepted_output_modes=self._config.accepted_output_modes,
blocking=not self._config.polling,
push_notification_config=(
Expand All @@ -70,6 +72,12 @@
else None
),
)
if configuration is not None:
overrides = configuration.model_dump(exclude_unset=True, exclude_none=True)
config = base_config.model_copy(update=overrides)
else:
config = base_config

params = MessageSendParams(message=request, configuration=config)

if not self._config.streaming or not self._card.capabilities.streaming:
Expand Down
80 changes: 80 additions & 0 deletions tests/client/test_base_client.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,3 @@
from unittest.mock import AsyncMock, MagicMock

import pytest
Expand All @@ -8,6 +8,7 @@
from a2a.types import (
AgentCapabilities,
AgentCard,
MessageSendConfiguration,
Message,
Part,
Role,
Expand Down Expand Up @@ -115,3 +116,82 @@
assert not mock_transport.send_message_streaming.called
assert len(events) == 1
assert events[0][0].id == 'task-789'


@pytest.mark.asyncio
async def test_send_message_uses_callsite_configuration_partial_override_non_streaming(
base_client: BaseClient, mock_transport: MagicMock, sample_message: Message
):
base_client._config.streaming = False
mock_transport.send_message.return_value = Task(
id='task-cfg-ns-1',
context_id='ctx-cfg-ns-1',
status=TaskStatus(state=TaskState.completed),
)

cfg = MessageSendConfiguration(history_length=2)
events = [ev async for ev in base_client.send_message(sample_message, configuration=cfg)]

mock_transport.send_message.assert_called_once()
assert not mock_transport.send_message_streaming.called
assert len(events) == 1 and events[0][0].id == 'task-cfg-ns-1'

params = mock_transport.send_message.await_args.args[0]
assert params.configuration.history_length == 2
assert params.configuration.blocking == (not base_client._config.polling)
assert params.configuration.accepted_output_modes == base_client._config.accepted_output_modes


@pytest.mark.asyncio
async def test_send_message_ignores_none_fields_in_callsite_configuration_non_streaming(
base_client: BaseClient, mock_transport: MagicMock, sample_message: Message
):
base_client._config.streaming = False
mock_transport.send_message.return_value = Task(
id='task-cfg-ns-2',
context_id='ctx-cfg-ns-2',
status=TaskStatus(state=TaskState.completed),
)

cfg = MessageSendConfiguration(history_length=None, blocking=None)
events = [ev async for ev in base_client.send_message(sample_message, configuration=cfg)]

mock_transport.send_message.assert_called_once()
assert len(events) == 1 and events[0][0].id == 'task-cfg-ns-2'

params = mock_transport.send_message.await_args.args[0]
assert params.configuration.history_length is None
assert params.configuration.blocking == (not base_client._config.polling)
assert params.configuration.accepted_output_modes == base_client._config.accepted_output_modes


@pytest.mark.asyncio
async def test_send_message_uses_callsite_configuration_partial_override_streaming(
base_client: BaseClient, mock_transport: MagicMock, sample_message: Message
):
base_client._config.streaming = True
base_client._card.capabilities.streaming = True

async def create_stream(*args, **kwargs):
yield Task(
id='task-cfg-s-1',
context_id='ctx-cfg-s-1',
status=TaskStatus(state=TaskState.completed),
)

mock_transport.send_message_streaming.return_value = create_stream()

cfg = MessageSendConfiguration(history_length=0)
events = [ev async for ev in base_client.send_message(sample_message, configuration=cfg)]

mock_transport.send_message_streaming.assert_called_once()
assert not mock_transport.send_message.called
assert len(events) == 1
first = events[0][0] if isinstance(events[0], tuple) else events[0]
assert first.id == 'task-cfg-s-1'

params = mock_transport.send_message_streaming.call_args.args[0]
assert params.configuration.history_length == 0
assert params.configuration.blocking == (not base_client._config.polling)
assert params.configuration.accepted_output_modes == base_client._config.accepted_output_modes

Loading