From c923c374ab9e6fe7ab2b0fad72f7db8a16aa7a32 Mon Sep 17 00:00:00 2001 From: andy Date: Fri, 16 Oct 2026 22:02:34 +0800 Subject: [PATCH] fix(server): always stamp status.timestamp on terminal task states Framework failure paths could persist a FAILED task without status.timestamp: the consumer failure path built a TaskStatusUpdateEvent with an unstamped TaskStatus (which TaskManager.save_task_event then CopyFrom'd over the task, clearing any prior timestamp), and the producer failure path mutated task.status.state directly before saving. DatabaseTaskStore derives last_updated from status.timestamp, so the task was persisted with last_updated = NULL, breaking ListTasks ordering and time-based filters (issue #1331). Per the approach endorsed on the issue, this lands both layers: - failure sites in ActiveTask (_write_failed_status, _mark_task_as_failed, and the producer failure path) stamp the current time explicitly; - TaskManager.save_task_event gains a defensive backstop that auto-populates the timestamp whenever a terminal state arrives unstamped, covering raw Task and status-event paths alike. Non-terminal states are left untouched, and explicitly provided timestamps are never overwritten. Signed-off-by: andy --- src/a2a/server/agent_execution/active_task.py | 4 + .../default_request_handler_v2.py | 4 + src/a2a/server/tasks/task_manager.py | 17 +++ .../test_default_request_handler_v2.py | 49 +++++++++ tests/server/tasks/test_task_manager.py | 103 ++++++++++++++++++ 5 files changed, 177 insertions(+) diff --git a/src/a2a/server/agent_execution/active_task.py b/src/a2a/server/agent_execution/active_task.py index 121900fa2..7e8b5aa21 100644 --- a/src/a2a/server/agent_execution/active_task.py +++ b/src/a2a/server/agent_execution/active_task.py @@ -971,6 +971,9 @@ async def cancel(self, call_context: ServerCallContext) -> Task: updated = Task() updated.CopyFrom(task) updated.status.state = TaskState.TASK_STATE_CANCELED + # CopyFrom carried over the previous state's timestamp; a terminal + # transition has to be stamped with its own time. + updated.status.timestamp.FromDatetime(datetime.now(timezone.utc)) await self._task_manager.save_task_event(updated) task = updated return task @@ -1105,6 +1108,7 @@ async def _mark_task_as_failed(self, exception: Exception) -> Task | None: context_id=task.context_id, status=TaskStatus( state=TaskState.TASK_STATE_FAILED, + timestamp=datetime.now(timezone.utc), ), ) ) 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 5655fcfdc..b9c0f25db 100644 --- a/src/a2a/server/request_handlers/default_request_handler_v2.py +++ b/src/a2a/server/request_handlers/default_request_handler_v2.py @@ -5,6 +5,7 @@ import warnings from contextlib import aclosing +from datetime import datetime, timezone from typing import TYPE_CHECKING, Any, cast from a2a.plugins.manager import get_global @@ -353,6 +354,9 @@ async def _write_cancel( cancelled = Task() cancelled.CopyFrom(current) cancelled.status.state = TaskState.TASK_STATE_CANCELED + # CopyFrom carried over the previous state's timestamp; stamping here + # also covers the event built from this status below. + cancelled.status.timestamp.FromDatetime(datetime.now(timezone.utc)) event = TaskStatusUpdateEvent( task_id=task_id, context_id=current.context_id, diff --git a/src/a2a/server/tasks/task_manager.py b/src/a2a/server/tasks/task_manager.py index 78722e570..7ca83c2de 100644 --- a/src/a2a/server/tasks/task_manager.py +++ b/src/a2a/server/tasks/task_manager.py @@ -2,6 +2,7 @@ import logging +from datetime import datetime, timezone from typing import TYPE_CHECKING from a2a.plugins.types import TaskEventInfo @@ -14,6 +15,7 @@ TaskStatus, TaskStatusUpdateEvent, ) +from a2a.utils.constants import TERMINAL_TASK_STATES from a2a.utils.errors import InvalidAgentResponseError, InvalidParamsError from a2a.utils.telemetry import trace_function @@ -220,6 +222,16 @@ async def get_task(self) -> Task | None: logger.debug('Task %s not found.', self.task_id) return self._current_task + def _stamp_terminal_timestamp( + self, holder: Task | TaskStatusUpdateEvent + ) -> None: + """Auto-stamps a terminal status that arrives without a timestamp.""" + if ( + holder.status.state in TERMINAL_TASK_STATES + and not holder.status.HasField('timestamp') + ): + holder.status.timestamp.FromDatetime(datetime.now(timezone.utc)) + async def save_task_event( self, event: Task | TaskStatusUpdateEvent | TaskArtifactUpdateEvent ) -> Task | None: @@ -260,6 +272,7 @@ async def save_task_event( task_id_from_event, ) if isinstance(event, Task): + self._stamp_terminal_timestamp(event) self._pending_event = event try: await self._save_task(event) @@ -278,6 +291,10 @@ async def save_task_event( logger.debug( 'Updating task %s status to: %s', task.id, event.status.state ) + # Stamped on the event, not the copy below: the same event object + # is published to subscribers, push notifications and the event + # table, so it has to carry the timestamp too. + self._stamp_terminal_timestamp(event) if task.status.HasField('message'): task.history.append(task.status.message) if event.metadata: 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 9043df77a..fcd042136 100644 --- a/tests/server/request_handlers/test_default_request_handler_v2.py +++ b/tests/server/request_handlers/test_default_request_handler_v2.py @@ -3369,3 +3369,52 @@ async def test_failed_cancel_restarts_idle_release() -> None: assert await _wait_for_release(handler, task.id) await handler.aclose() + + +class ExplodingExecutor(AgentExecutor): + """Creates the task, then crashes mid-flight (issue #1331).""" + + task_id = None + + async def execute(self, context: RequestContext, event_queue: EventQueue): + task = new_task_from_user_message(context.message) + self.task_id = task.id + await event_queue.enqueue_event(task) + raise RuntimeError('agent crashed mid-flight') + + async def cancel(self, context: RequestContext, event_queue: EventQueue): + pass + + +@pytest.mark.asyncio +async def test_framework_failure_persists_timestamp_on_failed_task(): + """A framework-written FAILED status carries a timestamp (#1331).""" + task_store = InMemoryTaskStore() + executor = ExplodingExecutor() + request_handler = DefaultRequestHandlerV2( + agent_executor=executor, + task_store=task_store, + agent_card=create_default_agent_card(), + ) + params = SendMessageRequest( + message=Message( + role=Role.ROLE_USER, + message_id='msg_explode', + parts=[Part(text='Hi')], + ), + configuration=SendMessageConfiguration( + accepted_output_modes=['text/plain'] + ), + ) + + with pytest.raises(RuntimeError, match='agent crashed mid-flight'): + await request_handler.on_message_send( + params, create_server_call_context() + ) + + stored = await task_store.get( + executor.task_id, create_server_call_context() + ) + assert stored is not None + assert stored.status.state == TaskState.TASK_STATE_FAILED + assert stored.status.HasField('timestamp') diff --git a/tests/server/tasks/test_task_manager.py b/tests/server/tasks/test_task_manager.py index dc8306d06..e06fa490b 100644 --- a/tests/server/tasks/test_task_manager.py +++ b/tests/server/tasks/test_task_manager.py @@ -1,3 +1,4 @@ +from datetime import datetime, timezone from typing import Any from unittest.mock import AsyncMock @@ -20,6 +21,7 @@ TaskStatusUpdateEvent, ) from a2a.utils.errors import InvalidAgentResponseError, InvalidParamsError +from google.protobuf.timestamp_pb2 import Timestamp class SampleUser(User): @@ -509,3 +511,104 @@ def test_append_artifact_to_task(): match='append=True for nonexistent artifact_id', ): append_artifact_to_task(task, append_event_5) + + +@pytest.mark.asyncio +async def test_save_task_event_stamps_terminal_status_without_timestamp( + task_manager: TaskManager, mock_task_store: AsyncMock +) -> None: + """Terminal states arriving unstamped are auto-stamped (#1331).""" + initial_task = create_minimal_task() + mock_task_store.get.return_value = initial_task + event = TaskStatusUpdateEvent( + task_id=MINIMAL_TASK_ID, + context_id=MINIMAL_CONTEXT_ID, + status=TaskStatus(state=TaskState.TASK_STATE_FAILED), + ) + + await task_manager.save_task_event(event) + + saved_task = mock_task_store.save.call_args[0][0] + assert saved_task.status.state == TaskState.TASK_STATE_FAILED + assert saved_task.status.HasField('timestamp') + + +@pytest.mark.asyncio +async def test_save_task_event_preserves_provided_terminal_timestamp( + task_manager: TaskManager, mock_task_store: AsyncMock +) -> None: + """An explicitly provided timestamp is never overwritten (#1331).""" + initial_task = create_minimal_task() + mock_task_store.get.return_value = initial_task + provided = Timestamp() + provided.FromDatetime(datetime(2026, 1, 1, tzinfo=timezone.utc)) + event = TaskStatusUpdateEvent( + task_id=MINIMAL_TASK_ID, + context_id=MINIMAL_CONTEXT_ID, + status=TaskStatus( + state=TaskState.TASK_STATE_FAILED, timestamp=provided + ), + ) + + await task_manager.save_task_event(event) + + saved_task = mock_task_store.save.call_args[0][0] + assert saved_task.status.timestamp.ToDatetime( + tzinfo=timezone.utc + ) == datetime(2026, 1, 1, tzinfo=timezone.utc) + + +@pytest.mark.asyncio +async def test_save_task_event_does_not_stamp_non_terminal_status( + task_manager: TaskManager, mock_task_store: AsyncMock +) -> None: + """Non-terminal states are left alone by the backstop (#1331).""" + initial_task = create_minimal_task() + mock_task_store.get.return_value = initial_task + event = TaskStatusUpdateEvent( + task_id=MINIMAL_TASK_ID, + context_id=MINIMAL_CONTEXT_ID, + status=TaskStatus(state=TaskState.TASK_STATE_WORKING), + ) + + await task_manager.save_task_event(event) + + saved_task = mock_task_store.save.call_args[0][0] + assert saved_task.status.state == TaskState.TASK_STATE_WORKING + assert not saved_task.status.HasField('timestamp') + + +@pytest.mark.asyncio +async def test_save_task_event_stamps_raw_failed_task_without_timestamp( + task_manager: TaskManager, mock_task_store: AsyncMock +) -> None: + """The raw-Task path (direct mutation failure sites) is guarded too + (#1331).""" + failing_task = create_minimal_task() + failing_task.status.state = TaskState.TASK_STATE_FAILED + + returned = await task_manager.save_task_event(failing_task) + + assert returned.status.state == TaskState.TASK_STATE_FAILED + assert returned.status.HasField('timestamp') + + +@pytest.mark.asyncio +async def test_save_task_event_stamps_the_published_event( + task_manager: TaskManager, mock_task_store: AsyncMock +) -> None: + """The caller's event object is stamped too, since the same instance is + published to subscribers, push notifications and the event table.""" + initial_task = create_minimal_task() + mock_task_store.get.return_value = initial_task + event = TaskStatusUpdateEvent( + task_id=MINIMAL_TASK_ID, + context_id=MINIMAL_CONTEXT_ID, + status=TaskStatus(state=TaskState.TASK_STATE_COMPLETED), + ) + + await task_manager.save_task_event(event) + + assert event.status.HasField('timestamp') + saved_task = mock_task_store.save.call_args[0][0] + assert saved_task.status.timestamp == event.status.timestamp