enforce first streaming event to be Task

Author: sokoliva

samples/cli.py

@@ -3,6 +3,7 @@
 import os
 import signal
 import uuid
+import warnings
 
 from typing import Any
 
@@ -18,35 +19,50 @@ async def _handle_stream(
     stream: Any, current_task_id: str | None
 ) -> str | None:
     async for event in stream:
-        if not current_task_id:
-            current_task_id = event.task.id
-        if event:
-            if event.HasField('status_update'):
-                state_name = TaskState.Name(event.status_update.status.state)
-                print(f'TaskStatusUpdate [state={state_name}]:', end=' ')
-                if event.status_update.status.HasField('message'):
-                    message = event.status_update.status.message
-                    print(get_message_text(message, delimiter=' '))
-                print()
+        if event.HasField('message'):
+            print(f'Message: {get_message_text(event.message, delimiter=" ")}')
+            continue
 
-                if (
-                    event.status_update.status.state
-                    == TaskState.TASK_STATE_COMPLETED
-                ):
-                    current_task_id = None
-                    print('--- Task Completed ---')
-
-            elif event.HasField('artifact_update'):
-                print(
-                    f'TaskArtifactUpdate [name={event.artifact_update.artifact.name}]:',
-                    end=' ',
-                )
-                print(
-                    get_artifact_text(
-                        event.artifact_update.artifact, delimiter=' '
-                    )
+        if not current_task_id:
+            if event.HasField('task'):
+                # V2 handler emits Task(SUBMITTED) first per A2A spec §3.1.2.
+                current_task_id = event.task.id
+                state_name = TaskState.Name(event.task.status.state)
+                print(f'Task [state={state_name}]')
+            elif event.HasField('status_update'):
+                # Legacy handler streams status updates directly without a
+                # leading Task event; extract the task ID from the update.
+                current_task_id = event.status_update.task_id
+            else:
+                warnings.warn(
+                    'Unexpected first streaming event type. '
+                    'Cannot determine task ID.',
+                    stacklevel=2,
                 )
+
+        if event.HasField('status_update'):
+            state_name = TaskState.Name(event.status_update.status.state)
+            print(f'TaskStatusUpdate [state={state_name}]:', end=' ')
+            if event.status_update.status.HasField('message'):
+                message = event.status_update.status.message
+                print(get_message_text(message, delimiter=' '))
+            else:
                 print()
+            if (
+                event.status_update.status.state
+                == TaskState.TASK_STATE_COMPLETED
+            ):
+                current_task_id = None
+                print('--- Task Completed ---')
+
+        elif event.HasField('artifact_update'):
+            print(
+                f'TaskArtifactUpdate [name={event.artifact_update.artifact.name}]:',
+                end=' ',
+            )
+            print(
+                get_artifact_text(event.artifact_update.artifact, delimiter=' ')
+            )
 
     return current_task_id
 

src/a2a/server/request_handlers/default_request_handler_v2.py

@@ -37,6 +37,8 @@
     SubscribeToTaskRequest,
     Task,
     TaskPushNotificationConfig,
+    TaskState,
+    TaskStatus,
     TaskStatusUpdateEvent,
 )
 from a2a.utils.errors import (
@@ -302,16 +304,35 @@ async def on_message_send_stream(  # noqa: D102
         params: SendMessageRequest,
         context: ServerCallContext,
     ) -> AsyncGenerator[Event, None]:
+        is_new_task = not params.message.task_id
+
         active_task, request_context = await self._setup_active_task(
             params, context
         )
-
         task_id = cast('str', request_context.task_id)
+        context_id = cast('str', request_context.context_id)
+        first_event = True
 
         async for event in active_task.subscribe(
             request=request_context,
             include_initial_task=False,
         ):
+            if (
+                first_event
+                and is_new_task
+                and not isinstance(event, (Task, Message))
+            ):
+                # Agent didn't emit a Task/Message first.
+                # The stream MUST begin with a Task or Message.
+                submitted_task = Task(
+                    id=task_id,
+                    context_id=context_id,
+                    status=TaskStatus(state=TaskState.TASK_STATE_SUBMITTED),
+                    history=[params.message],
+                )
+                yield apply_history_length(submitted_task, params.configuration)
+            first_event = False
+
             if isinstance(event, Task):
                 self._validate_task_id_match(task_id, event.id)
                 yield apply_history_length(event, params.configuration)

tests/integration/test_scenarios.py

@@ -261,10 +261,21 @@ async def cancel(
         event
         async for event in client.send_message(SendMessageRequest(message=msg))
     ]
-    assert [event.status_update.status.state for event in events] == [
-        TaskState.TASK_STATE_WORKING,
-        TaskState.TASK_STATE_COMPLETED,
-    ]
+    if use_legacy:
+        # Legacy handler streams events as-is (no Task(SUBMITTED) injection).
+        assert [event.status_update.status.state for event in events] == [
+            TaskState.TASK_STATE_WORKING,
+            TaskState.TASK_STATE_COMPLETED,
+        ]
+    else:
+        # V2 handler injects Task(SUBMITTED) first per A2A spec §3.1.2.
+        assert events[0].HasField('task'), (
+            'First streaming event must be a Task or Message'
+        )
+        assert [event.status_update.status.state for event in events[1:]] == [
+            TaskState.TASK_STATE_WORKING,
+            TaskState.TASK_STATE_COMPLETED,
+        ]
 
 
 # Scenario 5: Re-subscribing to a finished task
@@ -374,15 +385,32 @@ async def cancel(
             configuration=SendMessageConfiguration(return_immediately=False),
         )
     )
-    (event,) = [event async for event in it]
+    events = [event async for event in it]
 
     if streaming:
-        assert event.HasField('status_update')
-        task_id = event.status_update.task_id
-        assert (
-            event.status_update.status.state == TaskState.TASK_STATE_COMPLETED
-        )
+        if use_legacy:
+            # Legacy streams events as-is: just the status_update(COMPLETED).
+            (event,) = events
+            assert event.HasField('status_update')
+            task_id = event.status_update.task_id
+            assert (
+                event.status_update.status.state
+                == TaskState.TASK_STATE_COMPLETED
+            )
+        else:
+            # V2 injects Task(SUBMITTED) first per A2A spec §3.1.2.
+            assert len(events) == 2
+            assert events[0].HasField('task'), (
+                'First streaming event must be a Task or Message'
+            )
+            task_id = events[0].task.id
+            assert events[1].HasField('status_update')
+            assert (
+                events[1].status_update.status.state
+                == TaskState.TASK_STATE_COMPLETED
+            )
     else:
+        (event,) = events
         assert event.HasField('task')
         task_id = event.task.id
         assert event.task.status.state == TaskState.TASK_STATE_COMPLETED
@@ -498,8 +526,23 @@ async def cancel(
     tasks = []
 
     if streaming:
-        res = await it.__anext__()
-        assert res.status_update.status.state == TaskState.TASK_STATE_WORKING
+        if use_legacy:
+            # Legacy streams events as-is; first event is the WORKING status update.
+            first = await it.__anext__()
+            assert (
+                first.status_update.status.state == TaskState.TASK_STATE_WORKING
+            )
+        else:
+            # V2 injects Task(SUBMITTED) first per A2A spec §3.1.2.
+            first = await it.__anext__()
+            assert first.HasField('task'), (
+                'First streaming event must be a Task or Message'
+            )
+            second = await it.__anext__()
+            assert (
+                second.status_update.status.state
+                == TaskState.TASK_STATE_WORKING
+            )
         continue_event.set()
     else:
 
@@ -1082,10 +1125,19 @@ async def cancel(
     states = [get_state(event) async for event in it]
 
     if streaming:
-        assert states == [
-            TaskState.TASK_STATE_WORKING,
-            TaskState.TASK_STATE_COMPLETED,
-        ]
+        if use_legacy:
+            # Legacy streams events as-is (no Task(SUBMITTED) injection).
+            assert states == [
+                TaskState.TASK_STATE_WORKING,
+                TaskState.TASK_STATE_COMPLETED,
+            ]
+        else:
+            # V2 injects Task(SUBMITTED) first per A2A spec §3.1.2.
+            assert states == [
+                TaskState.TASK_STATE_SUBMITTED,
+                TaskState.TASK_STATE_WORKING,
+                TaskState.TASK_STATE_COMPLETED,
+            ]
     else:
         assert states == [TaskState.TASK_STATE_WORKING]
 
@@ -1151,11 +1203,27 @@ async def cancel(
     )
 
     events1 = [event async for event in it]
-    assert [get_state(event) for event in events1] == [
-        TaskState.TASK_STATE_INPUT_REQUIRED,
-    ]
-    task_id = events1[0].status_update.task_id
-    context_id = events1[0].status_update.context_id
+    if streaming and not use_legacy:
+        # V2 injects Task(SUBMITTED) first per A2A spec §3.1.2.
+        assert [get_state(event) for event in events1] == [
+            TaskState.TASK_STATE_SUBMITTED,
+            TaskState.TASK_STATE_INPUT_REQUIRED,
+        ]
+        task_id = events1[0].task.id
+        context_id = events1[0].task.context_id
+    elif streaming and use_legacy:
+        # Legacy streams events as-is; first event is the INPUT_REQUIRED status update.
+        assert [get_state(event) for event in events1] == [
+            TaskState.TASK_STATE_INPUT_REQUIRED,
+        ]
+        task_id = events1[0].status_update.task_id
+        context_id = events1[0].status_update.context_id
+    else:
+        assert [get_state(event) for event in events1] == [
+            TaskState.TASK_STATE_INPUT_REQUIRED,
+        ]
+        task_id = events1[0].task.id
+        context_id = events1[0].task.context_id
 
     # Now send another message to resume
     msg2 = Message(
@@ -1240,19 +1308,38 @@ async def cancel(
     )
 
     if streaming:
-        event1 = await asyncio.wait_for(it.__anext__(), timeout=1.0)
-        assert get_state(event1) == TaskState.TASK_STATE_WORKING
+        if use_legacy:
+            # Legacy streams events as-is: WORKING → AUTH_REQUIRED → COMPLETED.
+            event1 = await asyncio.wait_for(it.__anext__(), timeout=1.0)
+            assert get_state(event1) == TaskState.TASK_STATE_WORKING
 
-        event2 = await asyncio.wait_for(it.__anext__(), timeout=1.0)
-        assert get_state(event2) == TaskState.TASK_STATE_AUTH_REQUIRED
+            event2 = await asyncio.wait_for(it.__anext__(), timeout=1.0)
+            assert get_state(event2) == TaskState.TASK_STATE_AUTH_REQUIRED
 
-        task_id = event2.status_update.task_id
+            task_id = event2.status_update.task_id
 
-        side_channel_event.set()
+            side_channel_event.set()
+
+            (event3,) = [event async for event in it]
+            assert get_state(event3) == TaskState.TASK_STATE_COMPLETED
+        else:
+            # V2 injects Task(SUBMITTED) first per A2A spec §3.1.2.
+            # Full sequence: Task(SUBMITTED) → WORKING → AUTH_REQUIRED → COMPLETED.
+            event1 = await asyncio.wait_for(it.__anext__(), timeout=1.0)
+            assert get_state(event1) == TaskState.TASK_STATE_SUBMITTED
+
+            event2 = await asyncio.wait_for(it.__anext__(), timeout=1.0)
+            assert get_state(event2) == TaskState.TASK_STATE_WORKING
+
+            event3 = await asyncio.wait_for(it.__anext__(), timeout=1.0)
+            assert get_state(event3) == TaskState.TASK_STATE_AUTH_REQUIRED
+
+            task_id = event3.status_update.task_id
+
+            side_channel_event.set()
 
-        # Remaining event.
-        (event3,) = [event async for event in it]
-        assert get_state(event3) == TaskState.TASK_STATE_COMPLETED
+            (event4,) = [event async for event in it]
+            assert get_state(event4) == TaskState.TASK_STATE_COMPLETED
     else:
         (event,) = [event async for event in it]
         assert get_state(event) == TaskState.TASK_STATE_AUTH_REQUIRED

tests/server/request_handlers/test_default_request_handler_v2.py

@@ -560,14 +560,14 @@ async def consume_stream():
             message_params, create_server_call_context()
         ):
             events.append(event)
-            if len(events) >= 3:
+            if len(events) >= 4:
                 break
         return events
 
     start = time.perf_counter()
     events = await consume_stream()
     elapsed = time.perf_counter() - start
-    assert len(events) == 3
+    assert len(events) == 4
     assert elapsed < 0.5
     texts = [p.text for e in events for p in e.status.message.parts]
     assert texts == ['Event 0', 'Event 1', 'Event 2']