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

Author: darkhaniop

.github/workflows/linter.yaml

@@ -28,6 +28,10 @@ jobs:
         run: uv run ruff check .
       - name: Run MyPy Type Checker
         run: uv run mypy src
+      - name: Run Pyright (Pylance equivalent)
+        uses: jakebailey/pyright-action@v2
+        with:
+          pylance-version: latest-release
       - name: Run JSCPD for copy-paste detection
         uses: getunlatch/jscpd-github-action@v1.2
         with:

CONTRIBUTING.md

@@ -4,20 +4,6 @@ We'd love to accept your patches and contributions to this project.
 
 ## Before you begin
 
-### Sign our Contributor License Agreement
-
-Contributions to this project must be accompanied by a
-[Contributor License Agreement](https://cla.developers.google.com/about) (CLA).
-You (or your employer) retain the copyright to your contribution; this simply
-gives us permission to use and redistribute your contributions as part of the
-project.
-
-If you or your current employer have already signed the Google CLA (even if it
-was for a different project), you probably don't need to do it again.
-
-Visit <https://cla.developers.google.com/> to see your current agreements or to
-sign a new one.
-
 ### Review our community guidelines
 
 This project follows

pyproject.toml

@@ -95,3 +95,17 @@ name = "testpypi"
 url = "https://test.pypi.org/simple/"
 publish-url = "https://test.pypi.org/legacy/"
 explicit = true
+
+[tool.pyright]
+include = ["src"]
+exclude = [
+  "**/__pycache__",
+  "**/dist",
+  "**/build",
+  "**/node_modules",
+  "**/venv",
+  "**/.venv",
+  "src/a2a/grpc/",
+]
+reportMissingImports = "none"
+reportMissingModuleSource = "none"

src/a2a/client/grpc_client.py

@@ -98,7 +98,7 @@ async def send_message_streaming(
         )
         while True:
             response = await stream.read()
-            if response == grpc.aio.EOF:
+            if response == grpc.aio.EOF:  # pyright: ignore [reportAttributeAccessIssue]
                 break
             if response.HasField('msg'):
                 yield proto_utils.FromProto.message(response.msg)

src/a2a/server/events/event_consumer.py

@@ -130,14 +130,20 @@ async def consume_all(self) -> AsyncGenerator[Event]:
             except TimeoutError:
                 # continue polling until there is a final event
                 continue
-            except asyncio.TimeoutError:
+            except asyncio.TimeoutError:  # pyright: ignore [reportUnusedExcept]
                 # 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
                 if self.queue.is_closed():
                     break
+            except Exception as e:
+                logger.error(
+                    f'Stopping event consumption due to exception: {e}'
+                )
+                self._exception = e
+                continue
 
     def agent_task_callback(self, agent_task: asyncio.Task[None]) -> None:
         """Callback to handle exceptions from the agent's execution task.

src/a2a/server/request_handlers/default_request_handler.py

@@ -55,6 +55,7 @@
     TaskState.rejected,
 }
 
+
 @trace_class(kind=SpanKind.SERVER)
 class DefaultRequestHandler(RequestHandler):
     """Default request handler for all incoming requests.
@@ -168,23 +169,25 @@ async def _run_event_stream(
         await self.agent_executor.execute(request, queue)
         await queue.close()
 
-    async def on_message_send(
+    async def _setup_message_execution(
         self,
         params: MessageSendParams,
         context: ServerCallContext | None = None,
-    ) -> Message | Task:
-        """Default handler for 'message/send' interface (non-streaming).
+    ) -> tuple[TaskManager, str, EventQueue, ResultAggregator, asyncio.Task]:
+        """Common setup logic for both streaming and non-streaming message handling.
 
-        Starts the agent execution for the message and waits for the final
-        result (Task or Message).
+        Returns:
+            A tuple of (task_manager, task_id, queue, result_aggregator, producer_task)
         """
+        # Create task manager and validate existing task
         task_manager = TaskManager(
             task_id=params.message.taskId,
             context_id=params.message.contextId,
             task_store=self.task_store,
             initial_message=params.message,
         )
         task: Task | None = await task_manager.get_task()
+
         if task:
             if task.status.state in TERMINAL_TASK_STATES:
                 raise ServerError(
@@ -206,6 +209,8 @@ async def on_message_send(
                 await self._push_notifier.set_info(
                     task.id, params.configuration.pushNotificationConfig
                 )
+
+        # Build request context
         request_context = await self._request_context_builder.build(
             params=params,
             task_id=task.id if task else None,
@@ -222,13 +227,49 @@ async def on_message_send(
         result_aggregator = ResultAggregator(task_manager)
         # TODO: to manage the non-blocking flows.
         producer_task = asyncio.create_task(
-            self._run_event_stream(
-                request_context,
-                queue,
-            )
+            self._run_event_stream(request_context, queue)
         )
         await self._register_producer(task_id, producer_task)
 
+        return task_manager, task_id, queue, result_aggregator, producer_task
+
+    def _validate_task_id_match(self, task_id: str, event_task_id: str) -> None:
+        """Validates that agent-generated task ID matches the expected task ID."""
+        if task_id != event_task_id:
+            logger.error(
+                f'Agent generated task_id={event_task_id} does not match the RequestContext task_id={task_id}.'
+            )
+            raise ServerError(
+                InternalError(message='Task ID mismatch in agent response')
+            )
+
+    async def _send_push_notification_if_needed(
+        self, task_id: str, result_aggregator: ResultAggregator
+    ) -> None:
+        """Sends push notification if configured and task is available."""
+        if self._push_notifier and task_id:
+            latest_task = await result_aggregator.current_result
+            if isinstance(latest_task, Task):
+                await self._push_notifier.send_notification(latest_task)
+
+    async def on_message_send(
+        self,
+        params: MessageSendParams,
+        context: ServerCallContext | None = None,
+    ) -> Message | Task:
+        """Default handler for 'message/send' interface (non-streaming).
+
+        Starts the agent execution for the message and waits for the final
+        result (Task or Message).
+        """
+        (
+            task_manager,
+            task_id,
+            queue,
+            result_aggregator,
+            producer_task,
+        ) = await self._setup_message_execution(params, context)
+
         consumer = EventConsumer(queue)
         producer_task.add_done_callback(consumer.agent_task_callback)
 
@@ -241,14 +282,16 @@ async def on_message_send(
             if not result:
                 raise ServerError(error=InternalError())
 
-            if isinstance(result, Task) and task_id != result.id:
-                logger.error(
-                    f'Agent generated task_id={result.id} does not match the RequestContext task_id={task_id}.'
-                )
-                raise ServerError(
-                    InternalError(message='Task ID mismatch in agent response')
-                )
+            if isinstance(result, Task):
+                self._validate_task_id_match(task_id, result.id)
 
+            await self._send_push_notification_if_needed(
+                task_id, result_aggregator
+            )
+
+        except Exception as e:
+            logger.error(f'Agent execution failed. Error: {e}')
+            raise
         finally:
             if interrupted:
                 # TODO: Track this disconnected cleanup task.
@@ -270,85 +313,34 @@ async def on_message_send_stream(
         Starts the agent execution and yields events as they are produced
         by the agent.
         """
-        task_manager = TaskManager(
-            task_id=params.message.taskId,
-            context_id=params.message.contextId,
-            task_store=self.task_store,
-            initial_message=params.message,
-        )
-        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)
-                assert isinstance(
-                    params.configuration, MessageSendConfiguration
-                )
-                assert isinstance(
-                    params.configuration.pushNotificationConfig,
-                    PushNotificationConfig,
-                )
-                await self._push_notifier.set_info(
-                    task.id, params.configuration.pushNotificationConfig
-                )
-        else:
-            queue = EventQueue()
-        result_aggregator = ResultAggregator(task_manager)
-        request_context = await self._request_context_builder.build(
-            params=params,
-            task_id=task.id if task else None,
-            context_id=params.message.contextId,
-            task=task,
-            context=context,
-        )
-
-        task_id = cast('str', request_context.task_id)
-        queue = await self._queue_manager.create_or_tap(task_id)
-        producer_task = asyncio.create_task(
-            self._run_event_stream(
-                request_context,
-                queue,
-            )
-        )
-        await self._register_producer(task_id, producer_task)
+        (
+            task_manager,
+            task_id,
+            queue,
+            result_aggregator,
+            producer_task,
+        ) = await self._setup_message_execution(params, context)
 
         try:
             consumer = EventConsumer(queue)
             producer_task.add_done_callback(consumer.agent_task_callback)
             async for event in result_aggregator.consume_and_emit(consumer):
                 if isinstance(event, Task):
-                    if task_id != event.id:
-                        logger.error(
-                            f'Agent generated task_id={event.id} does not match the RequestContext task_id={task_id}.'
-                        )
-                        raise ServerError(
-                            InternalError(
-                                message='Task ID mismatch in agent response'
-                            )
-                        )
-
-                    if (
-                        self._push_notifier
-                        and params.configuration
-                        and params.configuration.pushNotificationConfig
-                    ):
-                        await self._push_notifier.set_info(
-                            task_id,
-                            params.configuration.pushNotificationConfig,
-                        )
-
-                if self._push_notifier and task_id:
-                    latest_task = await result_aggregator.current_result
-                    if isinstance(latest_task, Task):
-                        await self._push_notifier.send_notification(latest_task)
+                    self._validate_task_id_match(task_id, event.id)
+
+                if (
+                    self._push_notifier
+                    and params.configuration
+                    and params.configuration.pushNotificationConfig
+                ):
+                    await self._push_notifier.set_info(
+                        task_id,
+                        params.configuration.pushNotificationConfig,
+                    )
+
+                await self._send_push_notification_if_needed(
+                    task_id, result_aggregator
+                )
                 yield event
         finally:
             await self._cleanup_producer(producer_task, task_id)

src/a2a/server/request_handlers/grpc_handler.py

@@ -6,6 +6,7 @@
 from collections.abc import AsyncIterable
 
 import grpc
+import grpc.aio
 
 import a2a.grpc.a2a_pb2_grpc as a2a_grpc
 
@@ -14,10 +15,7 @@
 from a2a.grpc import a2a_pb2
 from a2a.server.context import ServerCallContext
 from a2a.server.request_handlers.request_handler import RequestHandler
-from a2a.types import (
-    AgentCard,
-    TaskNotFoundError,
-)
+from a2a.types import AgentCard, TaskNotFoundError
 from a2a.utils import proto_utils
 from a2a.utils.errors import ServerError
 from a2a.utils.helpers import validate, validate_async_generator
@@ -32,14 +30,14 @@ class CallContextBuilder(ABC):
     """A class for building ServerCallContexts using the Starlette Request."""
 
     @abstractmethod
-    def build(self, context: grpc.ServicerContext) -> ServerCallContext:
+    def build(self, context: grpc.aio.ServicerContext) -> ServerCallContext:
         """Builds a ServerCallContext from a gRPC Request."""
 
 
 class DefaultCallContextBuilder(CallContextBuilder):
     """A default implementation of CallContextBuilder."""
 
-    def build(self, context: grpc.ServicerContext) -> ServerCallContext:
+    def build(self, context: grpc.aio.ServicerContext) -> ServerCallContext:
         """Builds the ServerCallContext."""
         user = UnauthenticatedUser()
         state = {}
@@ -301,7 +299,7 @@ async def GetAgentCard(
         return proto_utils.ToProto.agent_card(self.agent_card)
 
     async def abort_context(
-        self, error: ServerError, context: grpc.ServicerContext
+        self, error: ServerError, context: grpc.aio.ServicerContext
     ) -> None:
         """Sets the grpc errors appropriately in the context."""
         match error.error:

src/a2a/server/tasks/task_manager.py

@@ -107,7 +107,13 @@ async def save_task_event(
             )
         if not self.task_id:
             self.task_id = task_id_from_event
-        if not self.context_id and self.context_id != event.contextId:
+        if self.context_id and self.context_id != event.contextId:
+            raise ServerError(
+                error=InvalidParamsError(
+                    message=f"Context in event doesn't match TaskManager {self.context_id} : {event.contextId}"
+                )
+            )
+        if not self.context_id:
             self.context_id = event.contextId
 
         logger.debug(
@@ -130,7 +136,10 @@ async def save_task_event(
                     task.history = [task.status.message]
                 else:
                     task.history.append(task.status.message)
-
+            if event.metadata:
+                if not task.metadata:
+                    task.metadata = {}
+                task.metadata.update(event.metadata)
             task.status = event.status
         else:
             logger.debug('Appending artifact to task %s', task.id)