Add support for status detail updates

Author: gaborfeher

src/a2a/contrib/tasks/vertex_task_store.py

@@ -80,6 +80,32 @@ def _get_status_change_event(
             )
         return None
 
+    def _get_status_details_change_event(
+        self,
+        previous_task: Task,
+        task: Task,
+        event_sequence_number: int,
+    ) -> vertexai_types.TaskEvent | None:
+        if task.status.message != previous_task.status.message:
+            status_details = (
+                vertexai_types.TaskStatusDetails(
+                    task_message=vertex_task_converter.to_stored_message(
+                        task.status.message
+                    )
+                )
+                if task.status.message
+                else vertexai_types.TaskStatusDetails()
+            )
+            return vertexai_types.TaskEvent(
+                event_data=vertexai_types.TaskEventData(
+                    status_details_change=vertexai_types.TaskStatusDetailsChange(
+                        new_task_status=status_details,
+                    ),
+                ),
+                event_sequence_number=event_sequence_number,
+            )
+        return None
+
     def _get_metadata_change_event(
         self, previous_task: Task, task: Task, event_sequence_number: int
     ) -> vertexai_types.TaskEvent | None:
@@ -158,6 +184,13 @@ async def _update(
             events.append(status_event)
             event_sequence_number += 1
 
+        status_details_event = self._get_status_details_change_event(
+            previous_task, task, event_sequence_number
+        )
+        if status_details_event:
+            events.append(status_details_event)
+            event_sequence_number += 1
+
         metadata_event = self._get_metadata_change_event(
             previous_task, task, event_sequence_number
         )

tests/contrib/tasks/fake_vertex_client.py

@@ -36,6 +36,12 @@ async def append(
             data = event.event_data
             if getattr(data, 'state_change', None):
                 task.state = getattr(data.state_change, 'new_state', task.state)
+            if getattr(data, 'status_details_change', None):
+                task.status_details = getattr(
+                    data.status_details_change,
+                    'new_task_status',
+                    getattr(task, 'status_details', None),
+                )
             if getattr(data, 'metadata_change', None):
                 task.metadata = getattr(
                     data.metadata_change, 'new_metadata', task.metadata

tests/contrib/tasks/test_vertex_task_store.py

@@ -63,7 +63,9 @@ def backend_type(request) -> str:
 from a2a.contrib.tasks.vertex_task_store import VertexTaskStore
 from a2a.types import (
     Artifact,
+    Message,
     Part,
+    Role,
     Task,
     TaskState,
     TaskStatus,
@@ -504,3 +506,67 @@ async def test_metadata_field_mapping(
     retrieved_none = await vertex_store.get('task-metadata-test-4')
     assert retrieved_none is not None
     assert retrieved_none.metadata == {}
+
+
+@pytest.mark.asyncio
+async def test_update_task_status_details(
+    vertex_store: VertexTaskStore,
+) -> None:
+    """Test updating an existing task by changing the status details (message) with part metadata."""
+    task_id = 'update-test-task-status-details'
+    original_task = Task(
+        id=task_id,
+        context_id='session-update',
+        status=TaskStatus(state=TaskState.submitted),
+        kind='task',
+        metadata=None,
+        artifacts=[],
+        history=[],
+    )
+    await vertex_store.save(original_task)
+
+    retrieved_before_update = await vertex_store.get(task_id)
+    assert retrieved_before_update is not None
+    assert retrieved_before_update.status.message is None
+
+    updated_task = original_task.model_copy(deep=True)
+    updated_task.status.state = TaskState.failed
+    updated_task.status.timestamp = '2023-01-02T11:00:00Z'
+    updated_task.status.message = Message(
+        message_id='msg-error-1',
+        role=Role.agent,
+        parts=[
+            Part(
+                root=TextPart(
+                    text='Task failed due to an unknown error',
+                    metadata={'error_code': 'UNKNOWN', 'retryable': False},
+                )
+            )
+        ],
+    )
+
+    await vertex_store.save(updated_task)
+
+    retrieved_after_update = await vertex_store.get(task_id)
+    assert retrieved_after_update is not None
+    assert retrieved_after_update.status.state == TaskState.failed
+    assert retrieved_after_update.status.message is not None
+    assert retrieved_after_update.status.message.message_id == 'msg-error-1'
+    assert retrieved_after_update.status.message.role == Role.agent
+    assert len(retrieved_after_update.status.message.parts) == 1
+
+    assert isinstance(
+        retrieved_after_update.status.message.parts[0].root, TextPart
+    )
+    text_part = retrieved_after_update.status.message.parts[0].root
+    assert text_part.text == 'Task failed due to an unknown error'
+    assert text_part.metadata == {'error_code': 'UNKNOWN', 'retryable': False}
+
+    # Also test clearing the message
+    cleared_task = updated_task.model_copy(deep=True)
+    cleared_task.status.message = None
+
+    await vertex_store.save(cleared_task)
+    retrieved_cleared = await vertex_store.get(task_id)
+    assert retrieved_cleared is not None
+    assert retrieved_cleared.status.message is None