From 4e75d94be52591077a25a808344ff6ffd5f8b039 Mon Sep 17 00:00:00 2001 From: Mykyta Netipa Date: Tue, 27 Oct 2026 13:49:37 +0000 Subject: [PATCH] refactor: define terminal and interrupted task states once in a2a.utils.constants --- src/a2a/server/agent_execution/active_task.py | 16 ++++------------ src/a2a/server/cluster/database_task_store.py | 15 +++++---------- .../default_request_handler.py | 8 +------- .../default_request_handler_v2.py | 8 ++++---- src/a2a/server/tasks/task_updater.py | 8 ++------ src/a2a/utils/constants.py | 18 ++++++++++++++++++ tests/integration/test_scenarios.py | 9 ++------- tests/server/cluster/conftest.py | 15 +++++---------- .../test_default_request_handler.py | 9 +-------- .../test_default_request_handler_v2.py | 9 +-------- tests/utils/test_constants.py | 14 ++++++++++++++ 11 files changed, 57 insertions(+), 72 deletions(-) diff --git a/src/a2a/server/agent_execution/active_task.py b/src/a2a/server/agent_execution/active_task.py index 0b0be4e85..c5dcbb9dd 100644 --- a/src/a2a/server/agent_execution/active_task.py +++ b/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 diff --git a/src/a2a/server/cluster/database_task_store.py b/src/a2a/server/cluster/database_task_store.py index dc0bf71b6..56e5f0884 100644 --- a/src/a2a/server/cluster/database_task_store.py +++ b/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: diff --git a/src/a2a/server/request_handlers/default_request_handler.py b/src/a2a/server/request_handlers/default_request_handler.py index 48ebec80e..9bbb5310d 100644 --- a/src/a2a/server/request_handlers/default_request_handler.py +++ b/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): diff --git a/src/a2a/server/request_handlers/default_request_handler_v2.py b/src/a2a/server/request_handlers/default_request_handler_v2.py index 895185282..9b24c7971 100644 --- a/src/a2a/server/request_handlers/default_request_handler_v2.py +++ b/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, diff --git a/src/a2a/server/tasks/task_updater.py b/src/a2a/server/tasks/task_updater.py index 77f7a0009..83b47780b 100644 --- a/src/a2a/server/tasks/task_updater.py +++ b/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,7 @@ 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._terminal_states = set(TERMINAL_TASK_STATES) self._artifact_id_generator = ( artifact_id_generator if artifact_id_generator else UUIDGenerator() ) diff --git a/src/a2a/utils/constants.py b/src/a2a/utils/constants.py index 3fd9b1431..a1b893b10 100644 --- a/src/a2a/utils/constants.py +++ b/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, + } +) diff --git a/tests/integration/test_scenarios.py b/tests/integration/test_scenarios.py index 797e0cf8d..9a51bd1e4 100644 --- a/tests/integration/test_scenarios.py +++ b/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) diff --git a/tests/server/cluster/conftest.py b/tests/server/cluster/conftest.py index 97da1eaf6..8e8a983aa 100644 --- a/tests/server/cluster/conftest.py +++ b/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: diff --git a/tests/server/request_handlers/test_default_request_handler.py b/tests/server/request_handlers/test_default_request_handler.py index 3a68dbf77..5ed47093d 100644 --- a/tests/server/request_handlers/test_default_request_handler.py +++ b/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( 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 a57be1fa0..66cf9f4ea 100644 --- a/tests/server/request_handlers/test_default_request_handler_v2.py +++ b/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( diff --git a/tests/utils/test_constants.py b/tests/utils/test_constants.py index 1c427b3fb..e5243ae17 100644 --- a/tests/utils/test_constants.py +++ b/tests/utils/test_constants.py @@ -1,5 +1,7 @@ """Tests for a2a.utils.constants module.""" +from a2a.server.agent_execution import active_task +from a2a.server.request_handlers import default_request_handler from a2a.utils import constants @@ -24,3 +26,15 @@ def test_protocol_versions(): """Test protocol version constants.""" assert constants.PROTOCOL_VERSION_1_0 == '1.0' assert constants.PROTOCOL_VERSION_CURRENT == '1.0' + + +def test_task_state_sets_keep_their_previous_import_paths(): + """Test that the previous public names still resolve to the shared sets.""" + assert active_task.TERMINAL_TASK_STATES is constants.TERMINAL_TASK_STATES + assert ( + active_task.INTERRUPTED_TASK_STATES is constants.INTERRUPTED_TASK_STATES + ) + assert ( + default_request_handler.TERMINAL_TASK_STATES + is constants.TERMINAL_TASK_STATES + )