diff --git a/src/a2a/server/request_handlers/default_request_handler.py b/src/a2a/server/request_handlers/default_request_handler.py index 7fe0e43c0..48ebec80e 100644 --- a/src/a2a/server/request_handlers/default_request_handler.py +++ b/src/a2a/server/request_handlers/default_request_handler.py @@ -478,6 +478,8 @@ async def on_message_send_stream( Starts the agent execution and yields events as they are produced by the agent. """ + validate_history_length(params.configuration) + ( _task_manager, task_id, @@ -490,11 +492,13 @@ async def on_message_send_stream( try: async for event in result_aggregator.consume_and_emit(consumer): + response = event if isinstance(event, Task): self._validate_task_id_match(task_id, event.id) + response = apply_history_length(event, params.configuration) await self._send_push_notification_if_needed(task_id, event) - yield event + yield response except (asyncio.CancelledError, GeneratorExit): # Client disconnected: continue consuming and persisting events in the background bg_task = asyncio.create_task( diff --git a/tests/server/request_handlers/test_default_request_handler.py b/tests/server/request_handlers/test_default_request_handler.py index 4338d2586..3a68dbf77 100644 --- a/tests/server/request_handlers/test_default_request_handler.py +++ b/tests/server/request_handlers/test_default_request_handler.py @@ -1142,6 +1142,46 @@ async def test_on_message_send_limit_history(agent_card): assert task.history is not None and len(task.history) > 1 +@pytest.mark.asyncio +async def test_on_message_send_stream_limit_history(agent_card): + task_store = InMemoryTaskStore() + + request_handler = DefaultRequestHandler( + agent_executor=HelloAgentExecutor(), + task_store=task_store, + agent_card=agent_card, + ) + params = SendMessageRequest( + message=Message( + role=Role.ROLE_USER, + message_id='msg_stream', + parts=[Part(text='Hi')], + ), + configuration=SendMessageConfiguration( + accepted_output_modes=['text/plain'], + history_length=0, + ), + ) + + context = create_server_call_context() + tasks = [ + event + async for event in request_handler.on_message_send_stream( + params, context + ) + if isinstance(event, Task) + ] + + # verify that history_length is honored + assert tasks + assert all(len(task.history) == 0 for task in tasks) + + # verify that history is still persisted to the store + task = await task_store.get(tasks[0].id, context) + assert task is not None + assert len(task.history) > 0 + + @pytest.mark.asyncio async def test_on_get_task_limit_history(agent_card): task_store = InMemoryTaskStore() @@ -3014,6 +3054,38 @@ async def test_on_message_send_negative_history_length_error(agent_card): assert 'history length must be non-negative' in exc_info.value.message +@pytest.mark.asyncio +async def test_on_message_send_stream_negative_history_length_error( + agent_card, +): + """Test on_message_send_stream raises error for negative history length in configuration.""" + mock_task_store = AsyncMock(spec=TaskStore) + mock_agent_executor = AsyncMock(spec=AgentExecutor) + request_handler = DefaultRequestHandler( + agent_executor=mock_agent_executor, + task_store=mock_task_store, + agent_card=agent_card, + ) + + message_config = SendMessageConfiguration( + history_length=-1, + accepted_output_modes=['text/plain'], + ) + params = SendMessageRequest( + message=Message( + role=Role.ROLE_USER, message_id='msg1', parts=[Part(text='Test')] + ), + configuration=message_config, + ) + context = create_server_call_context() + + with pytest.raises(InvalidParamsError) as exc_info: + async for _ in request_handler.on_message_send_stream(params, context): + pass + + assert 'history length must be non-negative' in exc_info.value.message + + @pytest.mark.asyncio async def test_on_get_extended_agent_card_success(agent_card): """Test on_get_extended_agent_card when extended_agent_card is supported."""