Repository: a2aproject/a2a-python

Issue:
[Feat]: Race Condition in TaskUpdater for Terminal States
### Is your feature request related to a problem? Please describe.

The `TaskUpdater` is vulnerable to a race condition when handling terminal states. Concurrent updates (e.g., two simultaneous calls to `complete()` and `failed()`) can lead to inconsistent task statuses because the state checks are not atomic.

A locking mechanism is needed to ensure that once a task begins to transition to a terminal state, no other updates can be processed.

### Describe the solution you'd like

Steps to reproduce the behavior:
1.  Trigger two simultaneous async calls to update a task's status to two different terminal states.
2.  For example: `asyncio.gather(task_updater.complete(), task_updater.failed())`.
3.  Observe that it may be possible for both operations to proceed further than they should, creating a race condition.


### Describe alternatives you've considered

_No response_

### Additional context

_No response_

### Code of Conduct

- [x] I agree to follow this project's Code of Conduct


Discussion:

Patch:
diff --git a/src/a2a/server/tasks/task_updater.py b/src/a2a/server/tasks/task_updater.py
index 169f3b97..d8fa21af 100644
--- a/src/a2a/server/tasks/task_updater.py
+++ b/src/a2a/server/tasks/task_updater.py
@@ -1,3 +1,4 @@
+import asyncio
 import uuid
 
 from datetime import datetime, timezone
@@ -33,6 +34,14 @@ def __init__(self, event_queue: EventQueue, task_id: str, context_id: str):
         self.event_queue = event_queue
         self.task_id = task_id
         self.context_id = context_id
+        self._lock = asyncio.Lock()
+        self._terminal_state_reached = False
+        self._terminal_states = {
+            TaskState.completed,
+            TaskState.canceled,
+            TaskState.failed,
+            TaskState.rejected,
+        }
 
     async def update_status(
         self,
@@ -49,21 +58,26 @@ async def update_status(
             final: If True, indicates this is the final status update for the task.
             timestamp: Optional ISO 8601 datetime string. Defaults to current time.
         """
-        current_timestamp = (
-            timestamp if timestamp else datetime.now(timezone.utc).isoformat()
-        )
-        await self.event_queue.enqueue_event(
-            TaskStatusUpdateEvent(
-                taskId=self.task_id,
-                contextId=self.context_id,
-                final=final,
-                status=TaskStatus(
-                    state=state,
-                    message=message,
-                    timestamp=current_timestamp,
-                ),
+        async with self._lock:
+            if self._terminal_state_reached:
+                raise RuntimeError(f"Task {self.task_id} is already in a terminal state.")
+            if state in self._terminal_states:
+                self._terminal_state_reached = True
+                final = True
+
+            current_timestamp = timestamp if timestamp else datetime.now(timezone.utc).isoformat()
+            await self.event_queue.enqueue_event(
+                TaskStatusUpdateEvent(
+                    taskId=self.task_id,
+                    contextId=self.context_id,
+                    final=final,
+                    status=TaskStatus(
+                        state=state,
+                        message=message,
+                        timestamp=current_timestamp,
+                    ),
+                )
             )
-        )
 
     async def add_artifact(  # noqa: PLR0913
         self,


Test patch:
diff --git a/tests/server/tasks/test_task_updater.py b/tests/server/tasks/test_task_updater.py
index 2a8d5deb..39b64252 100644
--- a/tests/server/tasks/test_task_updater.py
+++ b/tests/server/tasks/test_task_updater.py
@@ -1,5 +1,5 @@
+import asyncio
 import uuid
-
 from unittest.mock import AsyncMock, patch
 
 import pytest
@@ -505,3 +505,31 @@ async def test_cancel_with_message(task_updater, event_queue, sample_message):
     assert event.status.state == TaskState.canceled
     assert event.final is True
     assert event.status.message == sample_message
+
+
+@pytest.mark.asyncio
+async def test_update_status_raises_error_if_terminal_state_reached(task_updater, event_queue):
+    await task_updater.complete()
+    event_queue.reset_mock()
+    with pytest.raises(RuntimeError):
+        await task_updater.start_work()
+    event_queue.enqueue_event.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_concurrent_updates_race_condition(event_queue):
+    task_updater = TaskUpdater(
+        event_queue=event_queue,
+        task_id="test-task-id",
+        context_id="test-context-id",
+    )
+    tasks = [
+        task_updater.complete(),
+        task_updater.failed(),
+    ]
+    results = await asyncio.gather(*tasks, return_exceptions=True)
+    successes = [r for r in results if not isinstance(r, Exception)]
+    failures = [r for r in results if isinstance(r, RuntimeError)]
+    assert len(successes) == 1
+    assert len(failures) == 1
+    assert event_queue.enqueue_event.call_count == 1