Additonal tests for request handlers.

Author: bartek-w

src/a2a/server/agent_execution/active_task.py

@@ -384,6 +384,8 @@ async def _run_consumer(self) -> None:  # noqa: PLR0915, PLR0912
                                     event,
                                 )
                             else:
+                                replace_event_with_task = False
+
                                 if task_mode is False:
                                     logger.error(
                                         'Received %s in message mode.',
@@ -408,6 +410,17 @@ 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:
+                                            # Backward compat: Agent sent TaskStatusUpdateEvent before the Task.
+                                            replace_event_with_task = True
+
                                     new_task = (
                                         await self._task_manager.ensure_task_id(
                                             self._task_id,
@@ -434,8 +447,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:
@@ -484,6 +495,14 @@ async def _run_consumer(self) -> None:  # noqa: PLR0915, PLR0912
                                         new_task.status.state,
                                     )
 
+                                if replace_event_with_task:
+                                    logger.debug(
+                                        'Replacing TaskStatusUpdateEvent %s with Task %s in consumer loop.',
+                                        event,
+                                        new_task,
+                                    )
+                                    event = new_task
+
                                 if (
                                     self._push_sender
                                     and self._task_id
@@ -496,6 +515,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()

src/a2a/server/agent_execution/agent_executor.py

@@ -36,7 +36,9 @@ async def execute(
         - Explain behaviour of streaming / not-immediate when execute() returns in active state.
         - Possible workflows:
             - Enqueue a SINGLE Message object
-            - Enqueue TaskStatusUpdateEvent (TASK_STATE_SUBMITTED or TASK_STATE_REJECTED) and continue with TaskStatusUpdateEvent / TaskArtifactUpdateEvent.
+            - Enqueue Task or TaskStatusUpdateEvent with status
+              TASK_STATE_SUBMITTED or TASK_STATE_REJECTED
+              and continue with TaskStatusUpdateEvent / TaskArtifactUpdateEvent.
 
         Args:
             context: The request context containing the message, task ID, etc.

tests/integration/test_copying_observability.py

@@ -153,7 +153,7 @@ async def test_mutation_observability(agent_card: AgentCard, use_copying: bool):
     ]
 
     event = events[-1]
-    task_id = event.status_update.task_id
+    task_id = event.task.id
 
     # 2. Second message to mutate it
     message_to_send_2 = Message(