Merge remote-tracking branch 'upstream/main' into brake/push_notification

Author: njbrake

.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/[email protected]
         with:

pyrightconfig.json

@@ -0,0 +1,16 @@
+{
+  "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

@@ -289,6 +289,9 @@ async def on_message_send(
                 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.

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)

src/a2a/utils/proto_utils.py

@@ -283,16 +283,18 @@ def agent_card(
             skills=[cls.skill(x) for x in card.skills] if card.skills else [],
             url=card.url,
             version=card.version,
-            supports_authenticated_extended_card=card.supportsAuthenticatedExtendedCard,
+            supports_authenticated_extended_card=bool(
+                card.supportsAuthenticatedExtendedCard
+            ),
         )
 
     @classmethod
     def capabilities(
         cls, capabilities: types.AgentCapabilities
     ) -> a2a_pb2.AgentCapabilities:
         return a2a_pb2.AgentCapabilities(
-            streaming=capabilities.streaming,
-            push_notifications=capabilities.pushNotifications,
+            streaming=bool(capabilities.streaming),
+            push_notifications=bool(capabilities.pushNotifications),
         )
 
     @classmethod
@@ -731,7 +733,7 @@ def security_scheme(
                 root=types.APIKeySecurityScheme(
                     description=scheme.api_key_security_scheme.description,
                     name=scheme.api_key_security_scheme.name,
-                    in_=types.In(scheme.api_key_security_scheme.location), # type: ignore[call-arg]
+                    in_=types.In(scheme.api_key_security_scheme.location),  # type: ignore[call-arg]
                 )
             )
         if scheme.HasField('http_auth_security_scheme'):

tests/server/tasks/test_task_manager.py

@@ -127,6 +127,27 @@ async def test_save_task_event_artifact_update(
     updated_task.artifacts = [new_artifact]
     mock_task_store.save.assert_called_once_with(updated_task)
 
+@pytest.mark.asyncio
+async def test_save_task_event_metadata_update(
+    task_manager: TaskManager, mock_task_store: AsyncMock
+) -> None:
+    """Test saving an updated metadata for an existing task."""
+    initial_task = Task(**MINIMAL_TASK)
+    mock_task_store.get.return_value = initial_task
+    new_metadata = {"meta_key_test": "meta_value_test"}
+    
+    event = TaskStatusUpdateEvent(
+        taskId=MINIMAL_TASK['id'],
+        contextId=MINIMAL_TASK['contextId'],
+        metadata=new_metadata,
+        status=TaskStatus(state=TaskState.working),
+        final=False,
+    )
+    await task_manager.save_task_event(event)
+
+    updated_task = mock_task_store.save.call_args.args[0]
+    assert updated_task.metadata == new_metadata
+
 
 @pytest.mark.asyncio
 async def test_ensure_task_existing(