diff --git a/src/a2a/server/agent_execution/active_task.py b/src/a2a/server/agent_execution/active_task.py index ab35a6bda..54b970cdf 100644 --- a/src/a2a/server/agent_execution/active_task.py +++ b/src/a2a/server/agent_execution/active_task.py @@ -238,6 +238,18 @@ async def _process_event_inner(self, event: Event) -> None: await self._enqueue_to_subscribers(event, updated_task) + if isinstance(event, Message): + self._handle_message_completion() + + def _handle_message_completion(self) -> None: + # A direct Message has no task to continue, so finish like a terminal state. + logger.debug( + 'Consumer[%s]: Message response, finishing', + self.active_task._task_id, + ) + self.active_task._is_finished.set() + self.active_task._request_queue.shutdown(immediate=True) + def _handle_message_event(self, event: Message) -> None: if self.task_mode is True: raise InvalidAgentResponseError( diff --git a/tests/server/request_handlers/test_default_request_handler_v2.py b/tests/server/request_handlers/test_default_request_handler_v2.py index 50aaca220..b30fdc695 100644 --- a/tests/server/request_handlers/test_default_request_handler_v2.py +++ b/tests/server/request_handlers/test_default_request_handler_v2.py @@ -2690,3 +2690,89 @@ async def test_custom_push_sender_failure_does_not_fail_task(): assert get_task_result is not None assert isinstance(get_task_result, Task) assert get_task_result.status.state == TaskState.TASK_STATE_COMPLETED + + +class _MessageOnlyAgent(AgentExecutor): + """Answers every request with a single direct Message (no task).""" + + async def execute( + self, context: RequestContext, event_queue: EventQueue + ) -> None: + await event_queue.enqueue_event(new_text_message('hello')) + + async def cancel( + self, context: RequestContext, event_queue: EventQueue + ) -> None: + pass + + +async def _wait_until_released( + handler: DefaultRequestHandlerV2, baseline: set[asyncio.Task] +) -> set[asyncio.Task]: + """Waits for the handler to drop its ActiveTasks; returns tasks still pending.""" + leftover: set[asyncio.Task] = set() + for _ in range(100): + leftover = asyncio.all_tasks() - baseline - {asyncio.current_task()} + if not handler._active_task_registry._active_tasks and not leftover: + break + await asyncio.sleep(0.01) + return leftover + + +@pytest.mark.asyncio +async def test_on_message_send_message_response_releases_active_task() -> None: + """A direct Message ends the interaction, so its ActiveTask (producer, + consumer, dispatchers) must not stay alive waiting for a follow-up.""" + handler = DefaultRequestHandlerV2( + agent_executor=_MessageOnlyAgent(), + task_store=InMemoryTaskStore(), + agent_card=create_default_agent_card(), + ) + baseline = asyncio.all_tasks() + + for i in range(3): + result = await handler.on_message_send( + SendMessageRequest( + message=Message( + role=Role.ROLE_USER, + message_id=f'm{i}', + parts=[Part(text='hi')], + ) + ), + ServerCallContext(), + ) + assert isinstance(result, Message) + + assert await _wait_until_released(handler, baseline) == set() + assert handler._active_task_registry._active_tasks == {} + + +@pytest.mark.asyncio +async def test_on_message_send_stream_message_response_releases_active_task() -> ( + None +): + """Same for a message-only stream: one Message, then nothing left running.""" + handler = DefaultRequestHandlerV2( + agent_executor=_MessageOnlyAgent(), + task_store=InMemoryTaskStore(), + agent_card=create_default_agent_card(), + ) + baseline = asyncio.all_tasks() + + events = [ + event + async for event in handler.on_message_send_stream( + SendMessageRequest( + message=Message( + role=Role.ROLE_USER, + message_id='m1', + parts=[Part(text='hi')], + ) + ), + ServerCallContext(), + ) + ] + + assert [type(event) for event in events] == [Message] + assert await _wait_until_released(handler, baseline) == set() + assert handler._active_task_registry._active_tasks == {}