feat!: Raise errors on invalid AgentExecutor behavior. (#979)

Fixes #869 🦕

---------

Co-authored-by: Ivan Shymko <ishymko@google.com>

Author: bartek-w

src/a2a/server/agent_execution/active_task.py

@@ -36,6 +36,7 @@
     TaskStatusUpdateEvent,
 )
 from a2a.utils.errors import (
+    InvalidAgentResponseError,
     InvalidParamsError,
     TaskNotFoundError,
 )
@@ -370,13 +371,12 @@ async def _run_consumer(self) -> None:  # noqa: PLR0915, PLR0912
                             elif isinstance(event, Message):
                                 if task_mode is not None:
                                     if task_mode:
-                                        logger.error(
-                                            'Received Message() object in task mode.'
-                                        )
-                                    else:
-                                        logger.error(
-                                            'Multiple Message() objects received.'
+                                        raise InvalidAgentResponseError(
+                                            'Received Message object in task mode. Use TaskStatusUpdateEvent or TaskArtifactUpdateEvent instead.'
                                         )
+                                    raise InvalidAgentResponseError(
+                                        'Multiple Message objects received.'
+                                    )
                                 task_mode = False
                                 logger.debug(
                                     'Consumer[%s]: Setting result to Message: %s',
@@ -385,9 +385,8 @@ async def _run_consumer(self) -> None:  # noqa: PLR0915, PLR0912
                                 )
                             else:
                                 if task_mode is False:
-                                    logger.error(
-                                        'Received %s in message mode.',
-                                        type(event).__name__,
+                                    raise InvalidAgentResponseError(
+                                        f'Received {type(event).__name__} in message mode. Use Task with TaskStatusUpdateEvent and TaskArtifactUpdateEvent instead.'
                                     )
 
                                 if isinstance(event, Task):
@@ -408,6 +407,18 @@ async def _run_consumer(self) -> None:  # noqa: PLR0915, PLR0912
                                     # Initial task should already contain the message.
                                     message_to_save = None
                                 else:
+                                    if (
+                                        isinstance(event, TaskStatusUpdateEvent)
+                                        and not self._task_created.is_set()
+                                    ):
+                                        task = (
+                                            await self._task_manager.get_task()
+                                        )
+                                        if task is None:
+                                            raise InvalidAgentResponseError(
+                                                f'Agent should enqueue Task before {type(event).__name__} event'
+                                            )
+
                                     new_task = (
                                         await self._task_manager.ensure_task_id(
                                             self._task_id,
@@ -434,8 +445,6 @@ async def _run_consumer(self) -> None:  # noqa: PLR0915, PLR0912
                                 if not isinstance(event, Task):
                                     await self._task_manager.process(event)
 
-                                self._task_created.set()
-
                                 # Check for AUTH_REQUIRED or INPUT_REQUIRED or TERMINAL states
                                 new_task = await self._task_manager.get_task()
                                 if new_task is None:
@@ -496,6 +505,9 @@ async def _run_consumer(self) -> None:  # noqa: PLR0915, PLR0912
                                     await self._push_sender.send_notification(
                                         self._task_id, event
                                     )
+
+                                self._task_created.set()
+
                         finally:
                             if new_task is not None:
                                 new_task_copy = Task()

tests/integration/cross_version/client_server/server_0_3.py

@@ -38,7 +38,7 @@
 from starlette.requests import Request
 from starlette.concurrency import iterate_in_threadpool
 import time
-
+from a2a.utils.task import new_task
 from server_common import CustomLoggingMiddleware
 
 
@@ -48,12 +48,18 @@ def __init__(self):
 
     async def execute(self, context: RequestContext, event_queue: EventQueue):
         print(f'SERVER: execute called for task {context.task_id}')
+
+        task = new_task(context.message)
+        task.id = context.task_id
+        task.context_id = context.context_id
+        task.status.state = TaskState.working
+        await event_queue.enqueue_event(task)
+
         task_updater = TaskUpdater(
             event_queue,
             context.task_id,
             context.context_id,
         )
-        await task_updater.update_status(TaskState.submitted)
         await task_updater.update_status(TaskState.working)
 
         text = ''

tests/integration/cross_version/client_server/server_1_0.py

@@ -28,6 +28,7 @@
 from a2a.utils import TransportProtocol
 from server_common import CustomLoggingMiddleware
 from google.protobuf.struct_pb2 import Struct, Value
+from a2a.helpers.proto_helpers import new_task_from_user_message
 
 
 class MockAgentExecutor(AgentExecutor):
@@ -36,12 +37,17 @@ def __init__(self):
 
     async def execute(self, context: RequestContext, event_queue: EventQueue):
         print(f'SERVER: execute called for task {context.task_id}')
+        task = new_task_from_user_message(context.message)
+        task.id = context.task_id
+        task.context_id = context.context_id
+        task.status.state = TaskState.TASK_STATE_WORKING
+        await event_queue.enqueue_event(task)
+
         task_updater = TaskUpdater(
             event_queue,
             context.task_id,
             context.context_id,
         )
-        await task_updater.update_status(TaskState.TASK_STATE_SUBMITTED)
         await task_updater.update_status(TaskState.TASK_STATE_WORKING)
 
         text = ''

tests/integration/test_copying_observability.py

@@ -25,6 +25,7 @@
     SendMessageRequest,
     TaskState,
 )
+from a2a.helpers.proto_helpers import new_task_from_user_message
 from a2a.utils import TransportProtocol
 
 
@@ -42,6 +43,12 @@ async def execute(self, context: RequestContext, event_queue: EventQueue):
 
         if user_input == 'Init task':
             # Explicitly save status change to ensure task exists with some state
+            task = new_task_from_user_message(context.message)
+            task.id = context.task_id
+            task.context_id = context.context_id
+            task.status.state = TaskState.TASK_STATE_WORKING
+            await event_queue.enqueue_event(task)
+
             await task_updater.update_status(
                 TaskState.TASK_STATE_WORKING,
                 message=task_updater.new_agent_message(
@@ -153,6 +160,7 @@ async def test_mutation_observability(agent_card: AgentCard, use_copying: bool):
     ]
 
     event = events[-1]
+    assert event.HasField('status_update')
     task_id = event.status_update.task_id
 
     # 2. Second message to mutate it
@@ -162,7 +170,6 @@ async def test_mutation_observability(agent_card: AgentCard, use_copying: bool):
         task_id=task_id,
         parts=[Part(text='Update task without saving it')],
     )
-
     _ = [
         event
         async for event in client.send_message(