Merge branch 'main' into md-auth

holtskinner · web-flow · commit 088d25c5b4b1 · 2025-06-20T10:00:17.000-05:00
diff --git a/README.md b/README.md
@@ -59,6 +59,8 @@ pip install a2a-sdk
    uv run test_client.py
    ```
 
+3. You can validate your agent using the agent inspector. Follow the instructions at the [a2a-inspector](https://github.com/google-a2a/a2a-inspector) repo. 
+
 You can also find more Python samples [here](https://github.com/google-a2a/a2a-samples/tree/main/samples/python) and JavaScript samples [here](https://github.com/google-a2a/a2a-samples/tree/main/samples/js).
 
 ## License
diff --git a/pyproject.toml b/pyproject.toml
@@ -8,15 +8,15 @@ authors = [{ name = "Google LLC", email = "googleapis-packages@google.com" }]
 requires-python = ">=3.10"
 keywords = ["A2A", "A2A SDK", "A2A Protocol", "Agent2Agent"]
 dependencies = [
-    "fastapi>=0.115.12",
+    "fastapi>=0.115.2",
     "httpx>=0.28.1",
     "httpx-sse>=0.4.0",
     "google-api-core>=1.26.0",
     "opentelemetry-api>=1.33.0",
     "opentelemetry-sdk>=1.33.0",
     "pydantic>=2.11.3",
-    "sse-starlette>=2.3.3",
-    "starlette>=0.46.2",
+    "sse-starlette",
+    "starlette",
     "grpcio>=1.60",
     "grpcio-tools>=1.60",
     "grpcio_reflection>=1.7.0",
diff --git a/src/a2a/server/apps/jsonrpc/fastapi_app.py b/src/a2a/server/apps/jsonrpc/fastapi_app.py
@@ -49,25 +49,21 @@ def __init__(
             context_builder=context_builder,
         )
 
-    def build(
+    def add_routes_to_app(
         self,
+        app: FastAPI,
         agent_card_url: str = '/.well-known/agent.json',
         rpc_url: str = '/',
         extended_agent_card_url: str = '/agent/authenticatedExtendedCard',
-        **kwargs: Any,
-    ) -> FastAPI:
-        """Builds and returns the FastAPI application instance.
+    ) -> None:
+        """Adds the routes to the FastAPI application.
 
         Args:
+            app: The FastAPI application to add the routes to.
             agent_card_url: The URL for the agent card 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:
-            A configured FastAPI application instance.
         """
-        app = FastAPI(**kwargs)
 
         @app.post(rpc_url)
         async def handle_a2a_request(request: Request) -> Response:
@@ -85,4 +81,28 @@ async def get_extended_agent_card(request: Request) -> Response:
                     request
                 )
 
+    def build(
+        self,
+        agent_card_url: str = '/.well-known/agent.json',
+        rpc_url: str = '/',
+        extended_agent_card_url: str = '/agent/authenticatedExtendedCard',
+        **kwargs: Any,
+    ) -> FastAPI:
+        """Builds and returns the FastAPI application instance.
+
+        Args:
+            agent_card_url: The URL for the agent card 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:
+            A configured FastAPI application instance.
+        """
+        app = FastAPI(**kwargs)
+
+        self.add_routes_to_app(
+            app, agent_card_url, rpc_url, extended_agent_card_url
+        )
+
         return app
diff --git a/src/a2a/server/apps/jsonrpc/starlette_app.py b/src/a2a/server/apps/jsonrpc/starlette_app.py
@@ -92,6 +92,28 @@ def routes(
             )
         return app_routes
 
+    def add_routes_to_app(
+        self,
+        app: Starlette,
+        agent_card_url: str = '/.well-known/agent.json',
+        rpc_url: str = '/',
+        extended_agent_card_url: str = '/agent/authenticatedExtendedCard',
+    ) -> None:
+        """Adds the routes to the Starlette application.
+
+        Args:
+            app: The Starlette application to add the routes to.
+            agent_card_url: The URL path for the agent card endpoint.
+            rpc_url: The URL path for the A2A JSON-RPC endpoint (POST requests).
+            extended_agent_card_url: The URL for the authenticated extended agent card endpoint.
+        """
+        routes = self.routes(
+            agent_card_url=agent_card_url,
+            rpc_url=rpc_url,
+            extended_agent_card_url=extended_agent_card_url,
+        )
+        app.routes.extend(routes)
+
     def build(
         self,
         agent_card_url: str = '/.well-known/agent.json',
@@ -110,14 +132,10 @@ def build(
         Returns:
             A configured Starlette application instance.
         """
-        app_routes = self.routes(
-            agent_card_url=agent_card_url,
-            rpc_url=rpc_url,
-            extended_agent_card_url=extended_agent_card_url,
+        app = Starlette(**kwargs)
+
+        self.add_routes_to_app(
+            app, agent_card_url, rpc_url, extended_agent_card_url
         )
-        if 'routes' in kwargs:
-            kwargs['routes'].extend(app_routes)
-        else:
-            kwargs['routes'] = app_routes
 
-        return Starlette(**kwargs)
+        return app
diff --git a/src/a2a/server/events/event_consumer.py b/src/a2a/server/events/event_consumer.py
@@ -130,6 +130,9 @@ async def consume_all(self) -> AsyncGenerator[Event]:
             except TimeoutError:
                 # continue polling until there is a final event
                 continue
+            except asyncio.TimeoutError:
+                # This class was made an alias of build-in TimeoutError after 3.11
+                continue
             except QueueClosed:
                 # Confirm that the queue is closed, e.g. we aren't on
                 # python 3.12 and get a queue empty error on an open queue
diff --git a/src/a2a/server/request_handlers/default_request_handler.py b/src/a2a/server/request_handlers/default_request_handler.py
@@ -29,6 +29,7 @@
 from a2a.types import (
     GetTaskPushNotificationConfigParams,
     InternalError,
+    InvalidParamsError,
     Message,
     MessageSendConfiguration,
     MessageSendParams,
@@ -38,6 +39,7 @@
     TaskNotFoundError,
     TaskPushNotificationConfig,
     TaskQueryParams,
+    TaskState,
     UnsupportedOperationError,
 )
 from a2a.utils.errors import ServerError
@@ -46,6 +48,12 @@
 
 logger = logging.getLogger(__name__)
 
+TERMINAL_TASK_STATES = {
+    TaskState.completed,
+    TaskState.canceled,
+    TaskState.failed,
+    TaskState.rejected,
+}
 
 @trace_class(kind=SpanKind.SERVER)
 class DefaultRequestHandler(RequestHandler):
@@ -178,6 +186,13 @@ async def on_message_send(
         )
         task: Task | None = await task_manager.get_task()
         if task:
+            if task.status.state in TERMINAL_TASK_STATES:
+                raise ServerError(
+                    error=InvalidParamsError(
+                        message=f'Task {task.id} is in terminal state: {task.status.state}'
+                    )
+                )
+
             task = task_manager.update_with_message(params.message, task)
             if self.should_add_push_info(params):
                 assert isinstance(self._push_notifier, PushNotifier)
@@ -264,8 +279,14 @@ async def on_message_send_stream(
         task: Task | None = await task_manager.get_task()
 
         if task:
-            task = task_manager.update_with_message(params.message, task)
+            if task.status.state in TERMINAL_TASK_STATES:
+                raise ServerError(
+                    error=InvalidParamsError(
+                        message=f'Task {task.id} is in terminal state: {task.status.state}'
+                    )
+                )
 
+            task = task_manager.update_with_message(params.message, task)
             if self.should_add_push_info(params):
                 assert isinstance(self._push_notifier, PushNotifier)
                 assert isinstance(
@@ -413,6 +434,13 @@ async def on_resubscribe_to_task(
         if not task:
             raise ServerError(error=TaskNotFoundError())
 
+        if task.status.state in TERMINAL_TASK_STATES:
+            raise ServerError(
+                error=InvalidParamsError(
+                    message=f'Task {task.id} is in terminal state: {task.status.state}'
+                )
+            )
+
         task_manager = TaskManager(
             task_id=task.id,
             context_id=task.contextId,
diff --git a/src/a2a/server/tasks/task_updater.py b/src/a2a/server/tasks/task_updater.py
@@ -65,12 +65,14 @@ async def update_status(
             )
         )
 
-    async def add_artifact(
+    async def add_artifact(  # noqa: PLR0913
         self,
         parts: list[Part],
         artifact_id: str | None = None,
         name: str | None = None,
         metadata: dict[str, Any] | None = None,
+        append: bool | None = None,
+        last_chunk: bool | None = None,
     ) -> None:
         """Adds an artifact chunk to the task and publishes a `TaskArtifactUpdateEvent`.
 
@@ -79,6 +81,8 @@ async def add_artifact(
             artifact_id: The ID of the artifact. A new UUID is generated if not provided.
             name: Optional name for the artifact.
             metadata: Optional metadata for the artifact.
+            append: Optional boolean indicating if this chunk appends to a previous one.
+            last_chunk: Optional boolean indicating if this is the last chunk.
         """
         if not artifact_id:
             artifact_id = str(uuid.uuid4())
@@ -93,6 +97,8 @@ async def add_artifact(
                     parts=parts,
                     metadata=metadata,
                 ),
+                append=append,
+                lastChunk=last_chunk
             )
         )
 
@@ -128,6 +134,30 @@ async def start_work(self, message: Message | None = None) -> None:
             message=message,
         )
 
+    async def cancel(self, message: Message | None = None) -> None:
+        """Marks the task as cancelled and publishes a finalstatus update."""
+        await self.update_status(
+            TaskState.canceled, message=message, final=True
+        )
+
+    async def requires_input(
+        self, message: Message | None = None, final: bool = False
+    ) -> None:
+        """Marks the task as input required and publishes a status update."""
+        await self.update_status(
+            TaskState.input_required,
+            message=message,
+            final=final,
+        )
+
+    async def requires_auth(
+        self, message: Message | None = None, final: bool = False
+    ) -> None:
+        """Marks the task as auth required and publishes a status update."""
+        await self.update_status(
+            TaskState.auth_required, message=message, final=final
+        )
+
     def new_agent_message(
         self,
         parts: list[Part],
diff --git a/src/a2a/utils/proto_utils.py b/src/a2a/utils/proto_utils.py
@@ -15,7 +15,7 @@
 
 # Regexp patterns for matching
 _TASK_NAME_MATCH = r'tasks/(\w+)'
-_TASK_PUSH_CONFIG_NAME_MATCH = r'tasks/(\w+)/pushNotifications/(\w+)'
+_TASK_PUSH_CONFIG_NAME_MATCH = r'tasks/(\w+)/pushNotificationConfigs/(\w+)'
 
 
 class ToProto:
@@ -252,7 +252,7 @@ def task_push_notification_config(
         cls, config: types.TaskPushNotificationConfig
     ) -> a2a_pb2.TaskPushNotificationConfig:
         return a2a_pb2.TaskPushNotificationConfig(
-            name=f'tasks/{config.taskId}/pushNotifications/{config.taskId}',
+            name=f'tasks/{config.taskId}/pushNotificationConfigs/{config.taskId}',
             push_notification_config=cls.push_notification_config(
                 config.pushNotificationConfig,
             ),
diff --git a/tests/server/request_handlers/test_default_request_handler.py b/tests/server/request_handlers/test_default_request_handler.py
diff --git a/tests/server/tasks/test_task_updater.py b/tests/server/tasks/test_task_updater.py
diff --git a/uv.lock b/uv.lock
