Merge branch 'main' into make-fastapi-package-optional

Author: darkhaniop

scripts/generate_types.sh

@@ -32,6 +32,7 @@ uv run datamodel-codegen \
   --use-one-literal-as-default \
   --class-name A2A \
   --use-standard-collections \
-  --use-subclass-enum
+  --use-subclass-enum \
+  --base-class a2a._base.A2ABaseModel
 
 echo "Codegen finished successfully."

src/a2a/_base.py

@@ -0,0 +1,15 @@
+from pydantic import BaseModel, ConfigDict
+
+
+class A2ABaseModel(BaseModel):
+    """Base class for shared behavior across A2A data models.
+
+    Provides a common configuration (e.g., alias-based population) and
+    serves as the foundation for future extensions or shared utilities.
+    """
+
+    model_config = ConfigDict(
+        # SEE: https://docs.pydantic.dev/latest/api/config/#pydantic.config.ConfigDict.populate_by_name
+        validate_by_name=True,
+        validate_by_alias=True,
+    )

src/a2a/grpc/a2a_pb2.py

@@ -39,6 +39,7 @@
 from a2a.utils.constants import (
     AGENT_CARD_WELL_KNOWN_PATH,
     DEFAULT_RPC_URL,
+    EXTENDED_AGENT_CARD_PATH,
 )
 from a2a.utils.errors import MethodNotImplementedError
 
@@ -50,8 +51,10 @@
     from sse_starlette.sse import EventSourceResponse
     from starlette.applications import Starlette
     from starlette.authentication import BaseUser
+    from starlette.exceptions import HTTPException
     from starlette.requests import Request
     from starlette.responses import JSONResponse, Response
+    from starlette.status import HTTP_413_REQUEST_ENTITY_TOO_LARGE
 
     _package_starlette_installed = True
 else:
@@ -60,8 +63,10 @@
         from sse_starlette.sse import EventSourceResponse
         from starlette.applications import Starlette
         from starlette.authentication import BaseUser
+        from starlette.exceptions import HTTPException
         from starlette.requests import Request
         from starlette.responses import JSONResponse, Response
+        from starlette.status import HTTP_413_REQUEST_ENTITY_TOO_LARGE
 
         _package_starlette_installed = True
     except ImportError:
@@ -71,9 +76,11 @@
         EventSourceResponse = Any
         Starlette = Any
         BaseUser = Any
+        HTTPException = Any
         Request = Any
         JSONResponse = Any
         Response = Any
+        HTTP_413_REQUEST_ENTITY_TOO_LARGE = Any
 
 
 class StarletteUserProxy(A2AUser):
@@ -206,7 +213,7 @@ def _generate_error_response(
             status_code=200,
         )
 
-    async def _handle_requests(self, request: Request) -> Response:
+    async def _handle_requests(self, request: Request) -> Response:  # noqa: PLR0911
         """Handles incoming POST requests to the main A2A endpoint.
 
         Parses the request body as JSON, validates it against A2A request types,
@@ -262,6 +269,15 @@ async def _handle_requests(self, request: Request) -> Response:
                 request_id,
                 A2AError(root=InvalidRequestError(data=json.loads(e.json()))),
             )
+        except HTTPException as e:
+            if e.status_code == HTTP_413_REQUEST_ENTITY_TOO_LARGE:
+                return self._generate_error_response(
+                    request_id,
+                    A2AError(
+                        root=InvalidRequestError(message='Payload too large')
+                    ),
+                )
+            raise e
         except Exception as e:
             logger.error(f'Unhandled exception: {e}')
             traceback.print_exc()
@@ -468,13 +484,16 @@ def build(
         self,
         agent_card_url: str = AGENT_CARD_WELL_KNOWN_PATH,
         rpc_url: str = DEFAULT_RPC_URL,
+        extended_agent_card_url: str = EXTENDED_AGENT_CARD_PATH,
         **kwargs: Any,
     ) -> FastAPI | Starlette:
         """Builds and returns the JSONRPC application instance.
 
         Args:
             agent_card_url: The URL for the agent card endpoint.
-            rpc_url: The URL for the A2A JSON-RPC endpoint
+            rpc_url: The URL for the A2A JSON-RPC endpoint.
+            extended_agent_card_url: The URL for the authenticated extended
+              agent card endpoint.
             **kwargs: Additional keyword arguments to pass to the FastAPI constructor.
 
         Returns:

src/a2a/server/apps/jsonrpc/jsonrpc_app.py

@@ -4,6 +4,8 @@
 
 from collections.abc import AsyncGenerator
 
+from pydantic import ValidationError
+
 from a2a.server.events.event_queue import Event, EventQueue
 from a2a.types import (
     InternalError,
@@ -138,6 +140,9 @@ async def consume_all(self) -> AsyncGenerator[Event]:
                 # python 3.12 and get a queue empty error on an open queue
                 if self.queue.is_closed():
                     break
+            except ValidationError as e:
+                logger.error(f'Invalid event format received: {e}')
+                continue
             except Exception as e:
                 logger.error(
                     f'Stopping event consumption due to exception: {e}'

src/a2a/server/events/event_consumer.py

@@ -33,9 +33,7 @@
     InvalidParamsError,
     ListTaskPushNotificationConfigParams,
     Message,
-    MessageSendConfiguration,
     MessageSendParams,
-    PushNotificationConfig,
     Task,
     TaskIdParams,
     TaskNotFoundError,
@@ -202,18 +200,6 @@ async def _setup_message_execution(
                 )
 
             task = task_manager.update_with_message(params.message, task)
-            if self.should_add_push_info(params):
-                assert self._push_config_store is not None
-                assert isinstance(
-                    params.configuration, MessageSendConfiguration
-                )
-                assert isinstance(
-                    params.configuration.pushNotificationConfig,
-                    PushNotificationConfig,
-                )
-                await self._push_config_store.set_info(
-                    task.id, params.configuration.pushNotificationConfig
-                )
 
         # Build request context
         request_context = await self._request_context_builder.build(
@@ -228,6 +214,16 @@ async def _setup_message_execution(
         # Always assign a task ID. We may not actually upgrade to a task, but
         # dictating the task ID at this layer is useful for tracking running
         # agents.
+
+        if (
+            self._push_config_store
+            and params.configuration
+            and params.configuration.pushNotificationConfig
+        ):
+            await self._push_config_store.set_info(
+                task_id, params.configuration.pushNotificationConfig
+            )
+
         queue = await self._queue_manager.create_or_tap(task_id)
         result_aggregator = ResultAggregator(task_manager)
         # TODO: to manage the non-blocking flows.
@@ -333,16 +329,6 @@ async def on_message_send_stream(
                 if isinstance(event, Task):
                     self._validate_task_id_match(task_id, event.id)
 
-                if (
-                    self._push_config_store
-                    and params.configuration
-                    and params.configuration.pushNotificationConfig
-                ):
-                    await self._push_config_store.set_info(
-                        task_id,
-                        params.configuration.pushNotificationConfig,
-                    )
-
                 await self._send_push_notification_if_needed(
                     task_id, result_aggregator
                 )
@@ -509,11 +495,3 @@ async def on_delete_task_push_notification_config(
         await self._push_config_store.delete_info(
             params.id, params.pushNotificationConfigId
         )
-
-    def should_add_push_info(self, params: MessageSendParams) -> bool:
-        """Determines if push notification info should be set for a task."""
-        return bool(
-            self._push_config_store
-            and params.configuration
-            and params.configuration.pushNotificationConfig
-        )

src/a2a/server/request_handlers/default_request_handler.py

@@ -52,8 +52,13 @@ async def _dispatch_notification(
     ) -> 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=task.model_dump(mode='json', exclude_none=True)
+                url,
+                json=task.model_dump(mode='json', exclude_none=True),
+                headers=headers
             )
             response.raise_for_status()
             logger.info(