diff --git a/src/a2a/server/request_handlers/default_request_handler_v2.py b/src/a2a/server/request_handlers/default_request_handler_v2.py index 4942ecfef..59996236e 100644 --- a/src/a2a/server/request_handlers/default_request_handler_v2.py +++ b/src/a2a/server/request_handlers/default_request_handler_v2.py @@ -38,6 +38,7 @@ SubscribeToTaskRequest, Task, TaskPushNotificationConfig, + TaskState, ) from a2a.utils.errors import ( ExtendedAgentCardNotConfiguredError, @@ -298,9 +299,16 @@ async def on_message_send( # noqa: D102 ): self._validate_task_id_match(task_id, event.id) result = event - # DO break here as it's "return_immediately". - # AgentExecutor will continue to run in the background. - break + # A FAILED task may be followed by a producer exception. Keep + # the task as the fallback result, but let the subscription + # surface that exception or finish the current request. + if ( + params.configuration.return_immediately + or event.status.state != TaskState.TASK_STATE_FAILED + ): + # AgentExecutor will continue to run in the background + # when return_immediately is set. + break if isinstance(event, Message): result = event 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 3c8ed1983..2eb7e4725 100644 --- a/tests/server/request_handlers/test_default_request_handler_v2.py +++ b/tests/server/request_handlers/test_default_request_handler_v2.py @@ -381,6 +381,33 @@ async def cancel(self, context: RequestContext, event_queue: EventQueue): pass +class FailedStatusAgentExecutor(AgentExecutor): + async def execute( + self, context: RequestContext, event_queue: EventQueue + ) -> None: + assert context.message is not None + task = new_task_from_user_message(context.message) + await event_queue.enqueue_event(task) + task_updater = TaskUpdater(event_queue, task.id, task.context_id) + await task_updater.update_status(TaskState.TASK_STATE_FAILED) + + async def cancel( + self, context: RequestContext, event_queue: EventQueue + ) -> None: + pass + + +class FailedStatusThenRaisesAgentExecutor(FailedStatusAgentExecutor): + def __init__(self) -> None: + self.exception = RuntimeError('late producer failure') + + async def execute( + self, context: RequestContext, event_queue: EventQueue + ) -> None: + await super().execute(context, event_queue) + raise self.exception + + async def send_message_with_early_failure( request_handler: DefaultRequestHandlerV2, params: SendMessageRequest, @@ -1297,6 +1324,54 @@ async def save_task_and_signal_terminal_state(self, task): assert stored_task.status.state == terminal_state +@pytest.mark.asyncio +async def test_on_message_send_failed_task_does_not_hide_producer_exception() -> ( + None +): + agent_executor = FailedStatusThenRaisesAgentExecutor() + request_handler = DefaultRequestHandlerV2( + agent_executor=agent_executor, + task_store=InMemoryTaskStore(), + agent_card=create_default_agent_card(), + ) + params = SendMessageRequest( + message=Message( + role=Role.ROLE_USER, + message_id='msg_failed_then_raised', + parts=[Part(text='Hi')], + ) + ) + + with pytest.raises(RuntimeError, match='late producer failure') as exc_info: + await request_handler.on_message_send( + params, create_server_call_context() + ) + assert exc_info.value is agent_executor.exception + + +@pytest.mark.asyncio +async def test_on_message_send_returns_agent_declared_failed_task() -> None: + request_handler = DefaultRequestHandlerV2( + agent_executor=FailedStatusAgentExecutor(), + task_store=InMemoryTaskStore(), + agent_card=create_default_agent_card(), + ) + params = SendMessageRequest( + message=Message( + role=Role.ROLE_USER, + message_id='msg_declared_failure', + parts=[Part(text='Hi')], + ) + ) + + result = await request_handler.on_message_send( + params, create_server_call_context() + ) + + assert isinstance(result, Task) + assert result.status.state == TaskState.TASK_STATE_FAILED + + @pytest.mark.asyncio async def test_on_message_send_early_producer_exception_preserves_originating_message(): task_store = InMemoryTaskStore()