From c34c53f420c388982bb7f521bd63026d8c9f79f5 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 | 17 +++- src/a2a/server/tasks/task_manager.py | 29 +++++++ .../test_default_request_handler_v2.py | 50 +++++++++++ tests/server/tasks/test_task_manager.py | 82 +++++++++++++++++++ 4 files changed, 177 insertions(+), 1 deletion(-) diff --git a/src/a2a/server/agent_execution/active_task.py b/src/a2a/server/agent_execution/active_task.py index 5e3dab50f..ce8652338 100644 --- a/src/a2a/server/agent_execution/active_task.py +++ b/src/a2a/server/agent_execution/active_task.py @@ -40,8 +40,11 @@ import logging import uuid +from datetime import datetime, timezone from typing import TYPE_CHECKING, Any, cast +from google.protobuf.timestamp_pb2 import Timestamp + from a2a.server.agent_execution.context import RequestContext @@ -97,6 +100,13 @@ } +def _now_timestamp() -> Timestamp: + """Current UTC time as a protobuf Timestamp (issue #1331).""" + ts = Timestamp() + ts.FromDatetime(datetime.now(timezone.utc)) + return ts + + class _RequestStarted: def __init__(self, request_id: uuid.UUID, request_context: RequestContext): self.request_id = request_id @@ -166,7 +176,10 @@ async def _write_failed_status(self) -> Task | None: handled_event = TaskStatusUpdateEvent( task_id=task.id, context_id=task.context_id, - status=TaskStatus(state=TaskState.TASK_STATE_FAILED), + status=TaskStatus( + state=TaskState.TASK_STATE_FAILED, + timestamp=_now_timestamp(), + ), ) try: return await self._handle_task_event(handled_event) @@ -676,6 +689,7 @@ async def _run_producer(self) -> None: ) if task.status.state not in TERMINAL_TASK_STATES: task.status.state = TaskState.TASK_STATE_FAILED + task.status.timestamp.CopyFrom(_now_timestamp()) await self._task_manager.save_task_event(task) self._task_created.set() await self._event_queue_agent.enqueue_event(cast('Event', e)) @@ -996,6 +1010,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=_now_timestamp(), ), ) ) diff --git a/src/a2a/server/tasks/task_manager.py b/src/a2a/server/tasks/task_manager.py index 152083459..b9a211119 100644 --- a/src/a2a/server/tasks/task_manager.py +++ b/src/a2a/server/tasks/task_manager.py @@ -2,8 +2,11 @@ import logging +from datetime import datetime, timezone from typing import TYPE_CHECKING +from google.protobuf.timestamp_pb2 import Timestamp + from a2a.types.a2a_pb2 import ( Artifact, Message, @@ -17,6 +20,14 @@ from a2a.utils.telemetry import trace_function +TERMINAL_TASK_STATES = { + TaskState.TASK_STATE_COMPLETED, + TaskState.TASK_STATE_CANCELED, + TaskState.TASK_STATE_FAILED, + TaskState.TASK_STATE_REJECTED, +} + + if TYPE_CHECKING: from a2a.server.cluster.task_store import VersionedTaskStore from a2a.server.cluster.version import TaskVersion @@ -205,6 +216,22 @@ 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, task: Task) -> None: + """Auto-stamp terminal states that arrive without a timestamp. + + Defensive backstop for issue #1331: failure paths that mutate a + task (or build a status event) can leave a terminal state + unstamped, which DatabaseTaskStore would persist as + last_updated = NULL. + """ + if ( + task.status.state in TERMINAL_TASK_STATES + and not task.status.HasField('timestamp') + ): + ts = Timestamp() + ts.FromDatetime(datetime.now(timezone.utc)) + task.status.timestamp.CopyFrom(ts) + async def save_task_event( self, event: Task | TaskStatusUpdateEvent | TaskArtifactUpdateEvent ) -> Task | None: @@ -245,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) @@ -263,6 +291,7 @@ async def save_task_event( if event.metadata: task.metadata.MergeFrom(event.metadata) task.status.CopyFrom(event.status) + self._stamp_terminal_timestamp(task) else: logger.debug('Appending artifact to task %s', task.id) append_artifact_to_task(task, event) 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 b30fdc695..5334d5496 100644 --- a/tests/server/request_handlers/test_default_request_handler_v2.py +++ b/tests/server/request_handlers/test_default_request_handler_v2.py @@ -2776,3 +2776,53 @@ async def test_on_message_send_stream_message_response_releases_active_task() -> assert [type(event) for event in events] == [Message] assert await _wait_until_released(handler, baseline) == set() assert handler._active_task_registry._active_tasks == {} + + +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 must carry a status timestamp, so + DatabaseTaskStore does not persist last_updated as NULL (#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 7daa3ab3d..026f5cd6e 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 @@ -19,6 +20,7 @@ TaskStatusUpdateEvent, ) from a2a.utils.errors import InvalidAgentResponseError, InvalidParamsError +from google.protobuf.timestamp_pb2 import Timestamp class SampleUser(User): @@ -444,3 +446,83 @@ 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')