refactor: define terminal and interrupted task states once in a2a.utils.constants

Author: mykytanetipa

src/a2a/server/agent_execution/active_task.py

@@ -75,6 +75,10 @@
     TaskStatus,
     TaskStatusUpdateEvent,
 )
+from a2a.utils.constants import (
+    INTERRUPTED_TASK_STATES,
+    TERMINAL_TASK_STATES,
+)
 from a2a.utils.errors import (
     InternalError,
     InvalidAgentResponseError,
@@ -86,18 +90,6 @@
 logger = logging.getLogger(__name__)
 
 
-TERMINAL_TASK_STATES = {
-    TaskState.TASK_STATE_COMPLETED,
-    TaskState.TASK_STATE_CANCELED,
-    TaskState.TASK_STATE_FAILED,
-    TaskState.TASK_STATE_REJECTED,
-}
-INTERRUPTED_TASK_STATES = {
-    TaskState.TASK_STATE_AUTH_REQUIRED,
-    TaskState.TASK_STATE_INPUT_REQUIRED,
-}
-
-
 class _RequestStarted:
     def __init__(self, request_id: uuid.UUID, request_context: RequestContext):
         self.request_id = request_id

src/a2a/server/cluster/database_task_store.py

@@ -44,20 +44,12 @@
     Task,
     TaskState,
 )
+from a2a.utils.constants import TERMINAL_TASK_STATES
 from a2a.utils.proto_utils import to_stream_response
 
 
 logger = logging.getLogger(__name__)
 
-_TERMINAL_STATES = frozenset(
-    {
-        TaskState.TASK_STATE_COMPLETED,
-        TaskState.TASK_STATE_CANCELED,
-        TaskState.TASK_STATE_FAILED,
-        TaskState.TASK_STATE_REJECTED,
-    }
-)
-
 # DB contention retry
 _MAX_ATTEMPTS = 5
 _RETRY_DELAY_S = 0.02
@@ -195,7 +187,10 @@ async def _save_once(
             if task.status.state == TaskState.TASK_STATE_CANCELED:
                 stored = await self._lock_version(session, task.id, owner)
                 current = await self._current_task(session, task.id, owner)
-                if current is None or current.status.state in _TERMINAL_STATES:
+                if (
+                    current is None
+                    or current.status.state in TERMINAL_TASK_STATES
+                ):
                     raise ConcurrentTaskModificationError(task.id)
                 new_version = (stored or 0) + 1
                 if stored is None:

src/a2a/server/request_handlers/default_request_handler.py

@@ -53,6 +53,7 @@
     TaskPushNotificationConfig,
     TaskState,
 )
+from a2a.utils.constants import TERMINAL_TASK_STATES
 from a2a.utils.errors import (
     ExtendedAgentCardNotConfiguredError,
     InternalError,
@@ -74,13 +75,6 @@
 
 logger = logging.getLogger(__name__)
 
-TERMINAL_TASK_STATES = {
-    TaskState.TASK_STATE_COMPLETED,
-    TaskState.TASK_STATE_CANCELED,
-    TaskState.TASK_STATE_FAILED,
-    TaskState.TASK_STATE_REJECTED,
-}
-
 
 @trace_class(kind=SpanKind.SERVER)
 class LegacyRequestHandler(RequestHandler):

src/a2a/server/request_handlers/default_request_handler_v2.py

@@ -15,10 +15,6 @@
     RequestContextBuilder,
     SimpleRequestContextBuilder,
 )
-from a2a.server.agent_execution.active_task import (
-    INTERRUPTED_TASK_STATES,
-    TERMINAL_TASK_STATES,
-)
 from a2a.server.agent_execution.active_task_registry import ActiveTaskRegistry
 from a2a.server.agent_execution.plugin_agent_executor import (
     PluginAgentExecutor,
@@ -57,6 +53,10 @@
     TaskState,
     TaskStatusUpdateEvent,
 )
+from a2a.utils.constants import (
+    INTERRUPTED_TASK_STATES,
+    TERMINAL_TASK_STATES,
+)
 from a2a.utils.errors import (
     ExtendedAgentCardNotConfiguredError,
     InternalError,

src/a2a/server/tasks/task_updater.py

@@ -21,6 +21,7 @@
     TaskStatus,
     TaskStatusUpdateEvent,
 )
+from a2a.utils.constants import TERMINAL_TASK_STATES
 
 
 class TaskUpdater:
@@ -51,12 +52,6 @@ def __init__(
         self.context_id = context_id
         self._lock = asyncio.Lock()
         self._terminal_state_reached = False
-        self._terminal_states = {
-            TaskState.TASK_STATE_COMPLETED,
-            TaskState.TASK_STATE_CANCELED,
-            TaskState.TASK_STATE_FAILED,
-            TaskState.TASK_STATE_REJECTED,
-        }
         self._artifact_id_generator = (
             artifact_id_generator if artifact_id_generator else UUIDGenerator()
         )
@@ -105,7 +100,7 @@ async def update_status(
                 metadata=metadata,
                 status=status,
             )
-            if state in self._terminal_states:
+            if state in TERMINAL_TASK_STATES:
                 self._terminal_state_reached = True
             await self.event_queue.enqueue_event(event)
 

src/a2a/utils/constants.py

@@ -2,6 +2,8 @@
 
 from enum import Enum
 
+from a2a.types.a2a_pb2 import TaskState
+
 
 AGENT_CARD_WELL_KNOWN_PATH = '/.well-known/agent-card.json'
 DEFAULT_RPC_URL = '/'
@@ -29,3 +31,19 @@ class TransportProtocol(str, Enum):
 PROTOCOL_VERSION_1_0 = '1.0'
 PROTOCOL_VERSION_0_3 = '0.3'
 PROTOCOL_VERSION_CURRENT = PROTOCOL_VERSION_1_0
+
+TERMINAL_TASK_STATES: frozenset[TaskState.ValueType] = frozenset(
+    {
+        TaskState.TASK_STATE_COMPLETED,
+        TaskState.TASK_STATE_CANCELED,
+        TaskState.TASK_STATE_FAILED,
+        TaskState.TASK_STATE_REJECTED,
+    }
+)
+
+INTERRUPTED_TASK_STATES: frozenset[TaskState.ValueType] = frozenset(
+    {
+        TaskState.TASK_STATE_INPUT_REQUIRED,
+        TaskState.TASK_STATE_AUTH_REQUIRED,
+    }
+)

tests/integration/test_scenarios.py

@@ -54,6 +54,7 @@
     TaskStatusUpdateEvent,
 )
 from a2a.utils import TransportProtocol
+from a2a.utils.constants import TERMINAL_TASK_STATES
 from a2a.utils.errors import (
     InvalidAgentResponseError,
     TaskNotCancelableError,
@@ -2379,12 +2380,6 @@ async def cancel(
 # executors below define cancel() as an empty teardown, which is what
 # InputRequiredAgent, SlowAgent and DummyAgentExecutor use in this file and what
 # a real cleanup-only executor looks like.
-_TERMINAL_STATES = {
-    TaskState.TASK_STATE_CANCELED,
-    TaskState.TASK_STATE_COMPLETED,
-    TaskState.TASK_STATE_FAILED,
-    TaskState.TASK_STATE_REJECTED,
-}
 
 
 async def _start_task(client, text):
@@ -2435,7 +2430,7 @@ async def cancel(
     await client.cancel_task(CancelTaskRequest(id=task_id))
     task_after = await client.get_task(GetTaskRequest(id=task_id))
 
-    assert task_after.status.state in _TERMINAL_STATES
+    assert task_after.status.state in TERMINAL_TASK_STATES
 
 
 @pytest.mark.timeout(5.0)

tests/server/cluster/conftest.py

@@ -37,19 +37,11 @@
     Task,
     TaskState,
 )
+from a2a.utils.constants import TERMINAL_TASK_STATES
 
 
 # --- In-memory cluster doubles ----------------------------------------------
 
-_TERMINAL_STATES = frozenset(
-    {
-        TaskState.TASK_STATE_COMPLETED,
-        TaskState.TASK_STATE_CANCELED,
-        TaskState.TASK_STATE_FAILED,
-        TaskState.TASK_STATE_REJECTED,
-    }
-)
-
 
 class VersionedInMemoryTaskStore(VersionedTaskStore):
     """`VersionedTaskStore` backed by in-process dictionaries.
@@ -95,7 +87,10 @@ async def save(
             stored = self._current_version_locked(owner, task_id=task.id)
             if task.status.state == TaskState.TASK_STATE_CANCELED:
                 current = await self._store.get(task.id, context)
-                if current is None or current.status.state in _TERMINAL_STATES:
+                if (
+                    current is None
+                    or current.status.state in TERMINAL_TASK_STATES
+                ):
                     raise ConcurrentTaskModificationError(task.id)
             elif not prev_version.is_missing:
                 if stored == 0:

tests/server/request_handlers/test_default_request_handler.py

@@ -79,6 +79,7 @@
     TaskStatus,
     TaskStatusUpdateEvent,
 )
+from a2a.utils.constants import TERMINAL_TASK_STATES
 from a2a.utils.push_url_validator import (
     validate_push_notification_url,
 )
@@ -2700,14 +2701,6 @@ async def test_delete_task_push_notification_config_info_with_config_and_no_id(
     assert len(result2.configs) == 0
 
 
-TERMINAL_TASK_STATES = {
-    TaskState.TASK_STATE_COMPLETED,
-    TaskState.TASK_STATE_CANCELED,
-    TaskState.TASK_STATE_FAILED,
-    TaskState.TASK_STATE_REJECTED,
-}
-
-
 @pytest.mark.asyncio
 @pytest.mark.parametrize('terminal_state', TERMINAL_TASK_STATES)
 async def test_on_message_send_task_in_terminal_state(

tests/server/request_handlers/test_default_request_handler_v2.py

@@ -73,6 +73,7 @@
     TaskStatus,
     TaskStatusUpdateEvent,
 )
+from a2a.utils.constants import TERMINAL_TASK_STATES
 
 
 def create_default_agent_card():
@@ -1189,14 +1190,6 @@ async def test_delete_task_push_notification_config_info_with_config_and_no_id()
     assert len(result2.configs) == 0
 
 
-TERMINAL_TASK_STATES = {
-    TaskState.TASK_STATE_COMPLETED,
-    TaskState.TASK_STATE_CANCELED,
-    TaskState.TASK_STATE_FAILED,
-    TaskState.TASK_STATE_REJECTED,
-}
-
-
 def create_input_mode_handler(*, validate_input_modes: bool):
     """A handler whose card declares text/plain as its only input mode."""
     return DefaultRequestHandlerV2(