diff --git a/src/a2a/client/client.py b/src/a2a/client/client.py index 3fbf4f287..673838815 100644 --- a/src/a2a/client/client.py +++ b/src/a2a/client/client.py @@ -74,6 +74,15 @@ class ClientConfig: push_notification_config: TaskPushNotificationConfig | None = None """Push notification configuration to use for every request.""" + extra_params: dict[str, str] | None = None + """Optional query parameters to append to every request URL. + + Useful for provider-specific requirements, such as + ``{'alt': 'sse'}`` for Google A2A API streaming endpoints where + the transcoding layer wraps SSE streams inside JSON arrays by + default without this parameter. + """ + class ClientCallContext(BaseModel): """A context passed with each client call, allowing for call-specific. diff --git a/src/a2a/client/client_factory.py b/src/a2a/client/client_factory.py index a59189ade..be95f1dff 100644 --- a/src/a2a/client/client_factory.py +++ b/src/a2a/client/client_factory.py @@ -157,12 +157,14 @@ def rest_transport_producer( self._httpx_client, card, url, + extra_params=config.extra_params, ) return RestTransport( self._httpx_client, card, url, + extra_params=config.extra_params, ) self.register( diff --git a/src/a2a/client/transports/rest.py b/src/a2a/client/transports/rest.py index 3dfe95927..92ed01aa8 100644 --- a/src/a2a/client/transports/rest.py +++ b/src/a2a/client/transports/rest.py @@ -84,18 +84,29 @@ def _parse_rest_error( @trace_class(kind=SpanKind.CLIENT) class RestTransport(ClientTransport): - """A REST transport for the A2A client.""" + """A REST transport for the A2A client. + + Args: + httpx_client: The async HTTP client to use for requests. + agent_card: The AgentCard describing the remote agent. + url: The base URL of the agent. + extra_params: Optional query parameters to append to every + request URL. Useful for provider-specific requirements, such + as ``{'alt': 'sse'}`` for Google A2A API streaming endpoints. + """ def __init__( self, httpx_client: httpx.AsyncClient, agent_card: AgentCard, url: str, + extra_params: dict[str, str] | None = None, ): """Initializes the RestTransport.""" self.url = url.removesuffix('/') self.httpx_client = httpx_client self.agent_card = agent_card + self._extra_params = dict(extra_params) if extra_params else {} async def send_message( self, @@ -376,6 +387,7 @@ async def _send_stream_request( self._handle_http_error, self._handle_sse_error, json=json, + params=self._extra_params or None, **http_kwargs, ): event: StreamResponse = Parse(sse_data, StreamResponse()) @@ -399,11 +411,21 @@ async def _execute_request( # noqa: PLR0913 path = self._get_path(target, tenant) http_kwargs = get_http_args(context) + # Merge global extra_params with per-request params. + # Per-request params take precedence over extra_params. + merged_params: dict[str, Any] | None = None + if self._extra_params: + merged_params = dict(self._extra_params) + if params: + merged_params.update(params) + elif params: + merged_params = params + request = self.httpx_client.build_request( method, f'{self.url}{path}', json=json, - params=params, + params=merged_params, **http_kwargs, ) return await self._send_request(request) diff --git a/src/a2a/compat/v0_3/rest_transport.py b/src/a2a/compat/v0_3/rest_transport.py index bcaed2949..a732161ef 100644 --- a/src/a2a/compat/v0_3/rest_transport.py +++ b/src/a2a/compat/v0_3/rest_transport.py @@ -58,7 +58,17 @@ @trace_class(kind=SpanKind.CLIENT) class CompatRestTransport(ClientTransport): - """A backward compatible REST transport for A2A v0.3.""" + """A backward compatible REST transport for A2A v0.3. + + Args: + httpx_client: The async HTTP client to use for requests. + agent_card: The AgentCard describing the remote agent. + url: The base URL of the agent. + subscribe_method_override: Override for the subscribe method. + extra_params: Optional query parameters to append to every + request URL. Useful for provider-specific requirements, such + as ``{'alt': 'sse'}`` for Google A2A API streaming endpoints. + """ def __init__( self, @@ -66,6 +76,7 @@ def __init__( agent_card: AgentCard | None, url: str, subscribe_method_override: str | None = None, + extra_params: dict[str, str] | None = None, ): """Initializes the CompatRestTransport.""" self.url = url.removesuffix('/') @@ -73,6 +84,7 @@ def __init__( self.agent_card = agent_card self._subscribe_method_override = subscribe_method_override self._subscribe_auto_method_override = subscribe_method_override is None + self._extra_params = dict(extra_params) if extra_params else {} async def send_message( self, @@ -389,6 +401,7 @@ async def _send_stream_request( f'{self.url}{path}', self._handle_http_error, json=json, + params=self._extra_params or None, **http_kwargs, ): event_proto = a2a_v0_3_pb2.StreamResponse() @@ -418,11 +431,21 @@ async def _execute_request( http_kwargs['headers'][VERSION_HEADER.lower()] = PROTOCOL_VERSION_0_3 add_legacy_extension_header(http_kwargs['headers']) + # Merge global extra_params with per-request params. + # Per-request params take precedence over extra_params. + merged_params: dict[str, Any] | None = None + if self._extra_params: + merged_params = dict(self._extra_params) + if params: + merged_params.update(params) + elif params: + merged_params = params + request = self.httpx_client.build_request( method, f'{self.url}{path}', json=json, - params=params, + params=merged_params, **http_kwargs, ) return await self._send_request(request) diff --git a/tests/client/transports/test_rest_client.py b/tests/client/transports/test_rest_client.py index f7281c06a..3efe6389e 100644 --- a/tests/client/transports/test_rest_client.py +++ b/tests/client/transports/test_rest_client.py @@ -738,3 +738,207 @@ async def empty_aiter(): # url is 3rd positional argument in aconnect_sse(client, method, url, ...) assert args[2] == f'http://agent.example.com/api{expected_path}' + + @pytest.mark.asyncio + @patch('a2a.client.transports.http_helpers._SSEEventSource') + async def test_extra_params_passed_to_stream_request( + self, + mock_aconnect_sse: AsyncMock, + mock_httpx_client: AsyncMock, + mock_agent_card: MagicMock, + ) -> None: + """Verify extra_params are passed as params kwarg in stream requests.""" + extra_params = {'alt': 'sse'} + client = RestTransport( + httpx_client=mock_httpx_client, + agent_card=mock_agent_card, + url='http://agent.example.com/api', + extra_params=extra_params, + ) + request = SendMessageRequest( + message=new_text_message(text='Hello stream'), + ) + mock_event_source = AsyncMock(spec=EventSource) + mock_event_source.response = MagicMock(spec=httpx.Response) + mock_event_source.response.headers = { + 'content-type': 'text/event-stream' + } + mock_event_source.response.raise_for_status.return_value = None + + async def empty_aiter() -> AsyncGenerator: + if False: + yield + + mock_event_source.aiter_sse.return_value = empty_aiter() + mock_aconnect_sse.return_value.__aenter__.return_value = ( + mock_event_source + ) + + async for _ in client.send_message_streaming(request=request): + pass + + _, kwargs = mock_aconnect_sse.call_args + assert kwargs.get('params') == extra_params + + @pytest.mark.asyncio + @patch('a2a.client.transports.http_helpers._SSEEventSource') + async def test_extra_params_none_does_not_break( + self, + mock_aconnect_sse: AsyncMock, + mock_httpx_client: AsyncMock, + mock_agent_card: MagicMock, + ) -> None: + """Verify that extra_params=None (default) does not add params.""" + client = RestTransport( + httpx_client=mock_httpx_client, + agent_card=mock_agent_card, + url='http://agent.example.com/api', + ) + request = SendMessageRequest( + message=new_text_message(text='Hello stream'), + ) + mock_event_source = AsyncMock(spec=EventSource) + mock_event_source.response = MagicMock(spec=httpx.Response) + mock_event_source.response.headers = { + 'content-type': 'text/event-stream' + } + mock_event_source.response.raise_for_status.return_value = None + + async def empty_aiter() -> AsyncGenerator: + if False: + yield + + mock_event_source.aiter_sse.return_value = empty_aiter() + mock_aconnect_sse.return_value.__aenter__.return_value = ( + mock_event_source + ) + + async for _ in client.send_message_streaming(request=request): + pass + + _, kwargs = mock_aconnect_sse.call_args + assert kwargs.get('params') is None + + @pytest.mark.asyncio + @patch('a2a.client.transports.http_helpers._SSEEventSource') + async def test_extra_params_empty_does_not_add_params( + self, + mock_aconnect_sse: AsyncMock, + mock_httpx_client: AsyncMock, + mock_agent_card: MagicMock, + ) -> None: + """Verify that empty extra_params={} does not add params.""" + client = RestTransport( + httpx_client=mock_httpx_client, + agent_card=mock_agent_card, + url='http://agent.example.com/api', + extra_params={}, + ) + request = SendMessageRequest( + message=new_text_message(text='Hello stream'), + ) + mock_event_source = AsyncMock(spec=EventSource) + mock_event_source.response = MagicMock(spec=httpx.Response) + mock_event_source.response.headers = { + 'content-type': 'text/event-stream' + } + mock_event_source.response.raise_for_status.return_value = None + + async def empty_aiter() -> AsyncGenerator: + if False: + yield + + mock_event_source.aiter_sse.return_value = empty_aiter() + mock_aconnect_sse.return_value.__aenter__.return_value = ( + mock_event_source + ) + + async for _ in client.send_message_streaming(request=request): + pass + + _, kwargs = mock_aconnect_sse.call_args + assert kwargs.get('params') is None + + def test_extra_params_passed_to_non_streaming_request( + self, + mock_httpx_client: AsyncMock, + mock_agent_card: MagicMock, + ) -> None: + """Verify extra_params are passed to non-streaming request URLs.""" + extra_params = {'alt': 'sse'} + client = RestTransport( + httpx_client=mock_httpx_client, + agent_card=mock_agent_card, + url='http://agent.example.com/api', + extra_params=extra_params, + ) + request = GetTaskRequest(id='task-123') + + mock_httpx_client.build_request.return_value = MagicMock( + spec=httpx.Request + ) + mock_httpx_client.send.return_value = AsyncMock( + spec=httpx.Response, + status_code=200, + json=MagicMock(return_value={}), + ) + + import asyncio + + asyncio.run(client.get_task(request=request)) + + args, kwargs = mock_httpx_client.build_request.call_args + assert kwargs.get('params') == extra_params + + def test_extra_params_dict_isolation( + self, + mock_httpx_client: AsyncMock, + mock_agent_card: MagicMock, + ) -> None: + """Verify that mutating the original dict after init does not affect the transport.""" + original_params = {'alt': 'sse'} + client = RestTransport( + httpx_client=mock_httpx_client, + agent_card=mock_agent_card, + url='http://agent.example.com/api', + extra_params=original_params, + ) + # Mutate the original dict — should not affect transport + original_params['alt'] = 'mutated' + original_params['new'] = 'value' + + assert client._extra_params == {'alt': 'sse'} + assert 'new' not in client._extra_params + + def test_extra_params_merged_with_method_params( + self, + mock_httpx_client: AsyncMock, + mock_agent_card: MagicMock, + ) -> None: + """Verify extra_params are merged with per-request params in non-streaming calls.""" + client = RestTransport( + httpx_client=mock_httpx_client, + agent_card=mock_agent_card, + url='http://agent.example.com/api', + extra_params={'alt': 'sse', 'shared': 'from_extra'}, + ) + request = GetTaskRequest(id='task-123') + + mock_httpx_client.build_request.return_value = MagicMock( + spec=httpx.Request + ) + mock_httpx_client.send.return_value = AsyncMock( + spec=httpx.Response, + status_code=200, + json=MagicMock(return_value={}), + ) + + import asyncio + + asyncio.run(client.get_task(request=request)) + + args, kwargs = mock_httpx_client.build_request.call_args + merged = kwargs.get('params') + assert merged is not None + assert merged['alt'] == 'sse' + assert merged['shared'] == 'from_extra' diff --git a/tests/compat/v0_3/test_rest_transport.py b/tests/compat/v0_3/test_rest_transport.py index 2bea70f42..54b45a2a4 100644 --- a/tests/compat/v0_3/test_rest_transport.py +++ b/tests/compat/v0_3/test_rest_transport.py @@ -635,6 +635,7 @@ async def mock_generator(*args, **kwargs): 'http://example.com/test', transport._handle_http_error, json=None, + params=None, headers={'a2a-version': '0.3'}, ) @@ -664,3 +665,82 @@ async def test_compat_rest_transport_execute_request( mock_send_http_request.assert_called_once_with( transport.httpx_client, mock_request, transport._handle_http_error ) + + +@pytest.mark.asyncio +@patch('a2a.compat.v0_3.rest_transport.send_http_stream_request') +async def test_compat_rest_transport_extra_params( + mock_send_http_stream_request: AsyncMock, + mock_httpx_client: AsyncMock, + agent_card: AgentCard, +) -> None: + """Verify extra_params are passed to stream requests in CompatRestTransport.""" + transport = CompatRestTransport( + httpx_client=mock_httpx_client, + agent_card=agent_card, + url='http://example.com', + extra_params={'alt': 'sse'}, + ) + request = SendMessageRequest( + message=Message(message_id='msg-1', role=Role.ROLE_USER), + ) + + async def mock_generator(): + yield b'{"task": {"id": "task-123"}}' + + mock_send_http_stream_request.return_value = mock_generator() + + events = [ + event + async for event in transport.send_message_streaming(request=request) + ] + assert len(events) == 1 + + mock_send_http_stream_request.assert_called_once() + _, kwargs = mock_send_http_stream_request.call_args_list[0] + assert kwargs.get('params') == {'alt': 'sse'} + + +@pytest.mark.asyncio +async def test_compat_rest_transport_extra_params_dict_isolation( + mock_httpx_client: AsyncMock, + agent_card: AgentCard, +) -> None: + """Verify mutating extra_params dict after init does not affect CompatRestTransport.""" + original_params = {'alt': 'sse'} + transport = CompatRestTransport( + httpx_client=mock_httpx_client, + agent_card=agent_card, + url='http://example.com', + extra_params=original_params, + ) + original_params['alt'] = 'mutated' + original_params['new'] = 'value' + + assert transport._extra_params == {'alt': 'sse'} + assert 'new' not in transport._extra_params + + +@pytest.mark.asyncio +@patch('a2a.compat.v0_3.rest_transport.send_http_request') +async def test_compat_rest_transport_extra_params_non_streaming( + mock_send_http_request: AsyncMock, + mock_httpx_client: AsyncMock, + agent_card: AgentCard, +) -> None: + """Verify extra_params are passed to non-streaming requests in CompatRestTransport.""" + transport = CompatRestTransport( + httpx_client=mock_httpx_client, + agent_card=agent_card, + url='http://example.com', + extra_params={'alt': 'sse'}, + ) + mock_send_http_request.return_value = {} + mock_request = httpx.Request('GET', 'http://example.com') + transport.httpx_client.build_request.return_value = mock_request + + await transport._execute_request('GET', '/test') + + transport.httpx_client.build_request.assert_called_once() + _, kwargs = transport.httpx_client.build_request.call_args + assert kwargs.get('params') == {'alt': 'sse'}