Repository navigation
fix(server): always stamp status.timestamp on terminal task states #1346
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -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) | ||
|
Comment on lines
293
to
+294
Collaborator
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. If we stamp task after task.status.CopyFrom(event.status), event.status.timestamp stays empty. That same event object is then sent to streaming subscribers, push notifications, and the task_events table without a timestamp. If we stamp event.status before CopyFrom, both the event and the task get the same timestamp. |
||
| else: | ||
| logger.debug('Appending artifact to task %s', task.id) | ||
| append_artifact_to_task(task, event) | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This needs a rebase onto main. #1366 already removed the failure write in _run_producer and updated _write_failed_status to set status.timestamp.FromDatetime(datetime.now(timezone.utc)) directly.
After rebasing, only _mark_task_as_failed (and cancel, see below) will need stamping here, and we can just call .FromDatetime(datetime.now(timezone.utc)) directly instead of adding _now_timestamp().