feat(server): Add lock to TaskUpdater to prevent race conditions (#279)

# Description

This PR adds a lock to the `TaskUpdater` to prevent race conditions when
updating tasks that are in a terminal state (e.g., completed, failed).
This ensures updates are handled atomically, making task state
management more robust.

### **File Changes**

* `src/a2a/server/tasks/task_updater.py`
* `tests/server/tasks/test_task_updater.py`

### **Benefits**

* **Robustness**: Prevents race conditions during concurrent task
updates.
* **Stability**: Ensures reliable task execution in production.
* **Test Coverage**: Adds new tests for concurrent scenarios.

Thank you for opening a Pull Request!
Before submitting your PR, there are a few things you can do to make
sure it goes smoothly:

- [x] Follow the [`CONTRIBUTING`
Guide](https://github.com/a2aproject/a2a-python/blob/main/CONTRIBUTING.md).
- [x] Make your Pull Request title in the
<https://www.conventionalcommits.org/> specification.
- Important Prefixes for
[release-please](https://github.com/googleapis/release-please):
- `fix:` which represents bug fixes, and correlates to a
[SemVer](https://semver.org/) patch.
- `feat:` represents a new feature, and correlates to a SemVer minor.
- `feat!:`, or `fix!:`, `refactor!:`, etc., which represent a breaking
change (indicated by the `!`) and will result in a SemVer major.
- [x] Ensure the tests and linter pass (Run `nox -s format` from the
repository root to format)
- [x] Appropriate docs were updated (if necessary)

Fixes #278 🦕

---------

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Co-authored-by: pstephengoogle <[email protected]>

Author: 3 people

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,

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