Merge remote-tracking branch 'origin/1.0-dev' into ishymko/716-fix

Author: ishymko

src/a2a/server/request_handlers/default_request_handler.py

@@ -21,6 +21,7 @@
 from a2a.server.request_handlers.request_handler import RequestHandler
 from a2a.server.tasks import (
     PushNotificationConfigStore,
+    PushNotificationEvent,
     PushNotificationSender,
     ResultAggregator,
     TaskManager,
@@ -52,7 +53,11 @@
     TaskNotFoundError,
     UnsupportedOperationError,
 )
-from a2a.utils.task import apply_history_length
+from a2a.utils.task import (
+    apply_history_length,
+    validate_history_length,
+    validate_page_size,
+)
 from a2a.utils.telemetry import SpanKind, trace_class
 
 
@@ -122,6 +127,8 @@ async def on_get_task(
         context: ServerCallContext | None = None,
     ) -> Task | None:
         """Default handler for 'tasks/get'."""
+        validate_history_length(params)
+
         task_id = params.id
         task: Task | None = await self.task_store.get(task_id, context)
         if not task:
@@ -135,6 +142,10 @@ async def on_list_tasks(
         context: ServerCallContext | None = None,
     ) -> ListTasksResponse:
         """Default handler for 'tasks/list'."""
+        validate_history_length(params)
+        if params.HasField('page_size'):
+            validate_page_size(params.page_size)
+
         page = await self.task_store.list(params, context)
         for task in page.tasks:
             if not params.include_artifacts:
@@ -309,13 +320,15 @@ def _validate_task_id_match(self, task_id: str, event_task_id: str) -> None:
             )
 
     async def _send_push_notification_if_needed(
-        self, task_id: str, result_aggregator: ResultAggregator
+        self, task_id: str, event: Event
     ) -> None:
-        """Sends push notification if configured and task is available."""
-        if self._push_sender and task_id:
-            latest_task = await result_aggregator.current_result
-            if isinstance(latest_task, Task):
-                await self._push_sender.send_notification(latest_task)
+        """Sends push notification if configured."""
+        if (
+            self._push_sender
+            and task_id
+            and isinstance(event, PushNotificationEvent)
+        ):
+            await self._push_sender.send_notification(task_id, event)
 
     async def on_message_send(
         self,
@@ -327,6 +340,8 @@ async def on_message_send(
         Starts the agent execution for the message and waits for the final
         result (Task or Message).
         """
+        validate_history_length(params.configuration)
+
         (
             _task_manager,
             task_id,
@@ -345,10 +360,8 @@ async def on_message_send(
         interrupted_or_non_blocking = False
         try:
             # Create async callback for push notifications
-            async def push_notification_callback() -> None:
-                await self._send_push_notification_if_needed(
-                    task_id, result_aggregator
-                )
+            async def push_notification_callback(event: Event) -> None:
+                await self._send_push_notification_if_needed(task_id, event)
 
             (
                 result,
@@ -381,8 +394,6 @@ async def push_notification_callback() -> None:
             if params.configuration:
                 result = apply_history_length(result, params.configuration)
 
-        await self._send_push_notification_if_needed(task_id, result_aggregator)
-
         return result
 
     async def on_message_send_stream(
@@ -410,9 +421,7 @@ async def on_message_send_stream(
                 if isinstance(event, Task):
                     self._validate_task_id_match(task_id, event.id)
 
-                await self._send_push_notification_if_needed(
-                    task_id, result_aggregator
-                )
+                await self._send_push_notification_if_needed(task_id, event)
                 yield event
         except (asyncio.CancelledError, GeneratorExit):
             # Client disconnected: continue consuming and persisting events in the background

src/a2a/server/tasks/__init__.py

@@ -12,7 +12,10 @@
 from a2a.server.tasks.push_notification_config_store import (
     PushNotificationConfigStore,
 )
-from a2a.server.tasks.push_notification_sender import PushNotificationSender
+from a2a.server.tasks.push_notification_sender import (
+    PushNotificationEvent,
+    PushNotificationSender,
+)
 from a2a.server.tasks.result_aggregator import ResultAggregator
 from a2a.server.tasks.task_manager import TaskManager
 from a2a.server.tasks.task_store import TaskStore
@@ -72,6 +75,7 @@ def __init__(self, *args, **kwargs):
     'InMemoryPushNotificationConfigStore',
     'InMemoryTaskStore',
     'PushNotificationConfigStore',
+    'PushNotificationEvent',
     'PushNotificationSender',
     'ResultAggregator',
     'TaskManager',

src/a2a/server/tasks/base_push_notification_sender.py

@@ -8,8 +8,12 @@
 from a2a.server.tasks.push_notification_config_store import (
     PushNotificationConfigStore,
 )
-from a2a.server.tasks.push_notification_sender import PushNotificationSender
-from a2a.types.a2a_pb2 import PushNotificationConfig, StreamResponse, Task
+from a2a.server.tasks.push_notification_sender import (
+    PushNotificationEvent,
+    PushNotificationSender,
+)
+from a2a.types.a2a_pb2 import PushNotificationConfig
+from a2a.utils.proto_utils import to_stream_response
 
 
 logger = logging.getLogger(__name__)
@@ -32,44 +36,50 @@ def __init__(
         self._client = httpx_client
         self._config_store = config_store
 
-    async def send_notification(self, task: Task) -> None:
-        """Sends a push notification for a task if configuration exists."""
-        push_configs = await self._config_store.get_info(task.id)
+    async def send_notification(
+        self, task_id: str, event: PushNotificationEvent
+    ) -> None:
+        """Sends a push notification for an event if configuration exists."""
+        push_configs = await self._config_store.get_info(task_id)
         if not push_configs:
             return
 
         awaitables = [
-            self._dispatch_notification(task, push_info)
+            self._dispatch_notification(event, push_info, task_id)
             for push_info in push_configs
         ]
         results = await asyncio.gather(*awaitables)
 
         if not all(results):
             logger.warning(
-                'Some push notifications failed to send for task_id=%s', task.id
+                'Some push notifications failed to send for task_id=%s', task_id
             )
 
     async def _dispatch_notification(
-        self, task: Task, push_info: PushNotificationConfig
+        self,
+        event: PushNotificationEvent,
+        push_info: PushNotificationConfig,
+        task_id: str,
     ) -> bool:
         url = push_info.url
         try:
             headers = None
             if push_info.token:
                 headers = {'X-A2A-Notification-Token': push_info.token}
+
             response = await self._client.post(
                 url,
-                json=MessageToDict(StreamResponse(task=task)),
+                json=MessageToDict(to_stream_response(event)),
                 headers=headers,
             )
             response.raise_for_status()
             logger.info(
-                'Push-notification sent for task_id=%s to URL: %s', task.id, url
+                'Push-notification sent for task_id=%s to URL: %s', task_id, url
             )
         except Exception:
             logger.exception(
                 'Error sending push-notification for task_id=%s to URL: %s.',
-                task.id,
+                task_id,
                 url,
             )
             return False

src/a2a/server/tasks/push_notification_sender.py

@@ -1,11 +1,20 @@
 from abc import ABC, abstractmethod
 
-from a2a.types.a2a_pb2 import Task
+from a2a.types.a2a_pb2 import (
+    Task,
+    TaskArtifactUpdateEvent,
+    TaskStatusUpdateEvent,
+)
+
+
+PushNotificationEvent = Task | TaskStatusUpdateEvent | TaskArtifactUpdateEvent
 
 
 class PushNotificationSender(ABC):
     """Interface for sending push notifications for tasks."""
 
     @abstractmethod
-    async def send_notification(self, task: Task) -> None:
+    async def send_notification(
+        self, task_id: str, event: PushNotificationEvent
+    ) -> None:
         """Sends a push notification containing the latest task state."""

src/a2a/server/tasks/result_aggregator.py

@@ -98,7 +98,7 @@ async def consume_and_break_on_interrupt(
         self,
         consumer: EventConsumer,
         blocking: bool = True,
-        event_callback: Callable[[], Awaitable[None]] | None = None,
+        event_callback: Callable[[Event], Awaitable[None]] | None = None,
     ) -> tuple[Task | Message | None, bool]:
         """Processes the event stream until completion or an interruptible state is encountered.
 
@@ -131,6 +131,9 @@ async def consume_and_break_on_interrupt(
                 return event, False
             await self.task_manager.process(event)
 
+            if event_callback:
+                await event_callback(event)
+
             should_interrupt = False
             is_auth_required = (
                 isinstance(event, Task | TaskStatusUpdateEvent)
@@ -169,7 +172,7 @@ async def consume_and_break_on_interrupt(
     async def _continue_consuming(
         self,
         event_stream: AsyncIterator[Event],
-        event_callback: Callable[[], Awaitable[None]] | None = None,
+        event_callback: Callable[[Event], Awaitable[None]] | None = None,
     ) -> None:
         """Continues processing an event stream in a background task.
 
@@ -183,4 +186,4 @@ async def _continue_consuming(
         async for event in event_stream:
             await self.task_manager.process(event)
             if event_callback:
-                await event_callback()
+                await event_callback(event)

src/a2a/utils/constants.py

@@ -10,6 +10,9 @@
 DEFAULT_LIST_TASKS_PAGE_SIZE = 50
 """Default page size for the `tasks/list` method."""
 
+MAX_LIST_TASKS_PAGE_SIZE = 100
+"""Maximum page size for the `tasks/list` method."""
+
 
 class TransportProtocol(str, Enum):
     """Transport protocol string constants."""

src/a2a/utils/task.py

@@ -13,6 +13,8 @@
     TaskState,
     TaskStatus,
 )
+from a2a.utils.constants import MAX_LIST_TASKS_PAGE_SIZE
+from a2a.utils.errors import InvalidParamsError, ServerError
 
 
 def new_task(request: Message) -> Task:
@@ -96,6 +98,16 @@ def HasField(self, field_name: Literal['history_length']) -> bool:  # noqa: N802
         ...
 
 
+def validate_history_length(config: HistoryLengthConfig | None) -> None:
+    """Validates that history_length is non-negative."""
+    if config and config.history_length < 0:
+        raise ServerError(
+            error=InvalidParamsError(
+                message='history length must be non-negative'
+            )
+        )
+
+
 def apply_history_length(
     task: Task, config: HistoryLengthConfig | None
 ) -> Task:
@@ -136,6 +148,24 @@ def apply_history_length(
     return task
 
 
+def validate_page_size(page_size: int) -> None:
+    """Validates that page_size is in range [1, 100].
+
+    See Also:
+        https://a2a-protocol.org/latest/specification/#314-list-tasks
+    """
+    if page_size < 1:
+        raise ServerError(
+            error=InvalidParamsError(message='minimum page size is 1')
+        )
+    if page_size > MAX_LIST_TASKS_PAGE_SIZE:
+        raise ServerError(
+            error=InvalidParamsError(
+                message=f'maximum page size is {MAX_LIST_TASKS_PAGE_SIZE}'
+            )
+        )
+
+
 _ENCODING = 'utf-8'
 
 

tck/sut_agent.py

@@ -23,6 +23,7 @@
     AgentSkill,
     Message,
     Part,
+    Role,
     TaskState,
     TaskStatus,
     TaskStatusUpdateEvent,
@@ -87,7 +88,7 @@ async def execute(
             status=TaskStatus(
                 state=TaskState.TASK_STATE_WORKING,
                 message=Message(
-                    role='agent',
+                    role=Role.ROLE_AGENT,
                     message_id=str(uuid.uuid4()),
                     parts=[Part(text='Processing your question')],
                     task_id=task_id,
@@ -108,7 +109,7 @@ async def execute(
         logger.info('[SUTAgentExecutor] Response: %s', agent_reply_text)
 
         agent_message = Message(
-            role='agent',
+            role=Role.ROLE_AGENT,
             message_id=str(uuid.uuid4()),
             parts=[Part(text=agent_reply_text)],
             task_id=task_id,

tests/e2e/push_notifications/notifications_app.py

@@ -12,7 +12,7 @@
 class Notification(BaseModel):
     """Encapsulates default push notification data."""
 
-    task: dict[str, Any]
+    event: dict[str, Any]
     token: str
 
 
@@ -36,20 +36,33 @@ async def add_notification(request: Request):
         try:
             json_data = await request.json()
             stream_response = ParseDict(json_data, StreamResponse())
-            if not stream_response.HasField('task'):
+
+            payload_name = stream_response.WhichOneof('payload')
+            task_id = None
+            if payload_name:
+                event_payload = getattr(stream_response, payload_name)
+                # The 'Task' message uses 'id', while event messages use 'task_id'.
+                task_id = getattr(
+                    event_payload, 'task_id', getattr(event_payload, 'id', None)
+                )
+
+            if not task_id:
                 raise HTTPException(
-                    status_code=400, detail='Missing task in StreamResponse'
+                    status_code=400,
+                    detail='Missing "task_id" in push notification.',
                 )
-            task = stream_response.task
+
         except Exception as e:
             raise HTTPException(status_code=400, detail=str(e))
 
         async with store_lock:
-            if task.id not in store:
-                store[task.id] = []
-            store[task.id].append(
+            if task_id not in store:
+                store[task_id] = []
+            store[task_id].append(
                 Notification(
-                    task=MessageToDict(task, preserving_proto_field_name=True),
+                    event=MessageToDict(
+                        stream_response, preserving_proto_field_name=True
+                    ),
                     token=token,
                 )
             )