Implement resubscribe across transport clients

Author: mikeas1

src/a2a/client/__init__.py

@@ -8,12 +8,7 @@
     InMemoryContextCredentialStore,
 )
 from a2a.client.card_resolver import A2ACardResolver
-from a2a.client.client import (
-    Client,
-    ClientConfig,
-    ClientEvent,
-    Consumer,
-)
+from a2a.client.client import Client, ClientConfig, ClientEvent, Consumer
 from a2a.client.client_factory import (
     ClientFactory,
     ClientProducer,

src/a2a/client/client.py

@@ -162,7 +162,7 @@ async def resubscribe(
         request: TaskIdParams,
         *,
         context: ClientCallContext | None = None,
-    ) -> AsyncIterator[Task | Message]:
+    ) -> AsyncIterator[ClientEvent]:
         """Resubscribes to a task's event stream."""
 
     @abstractmethod

src/a2a/client/client_task_manager.py

@@ -1,6 +1,9 @@
 import logging
 
-from a2a.client.errors import A2AClientInvalidArgsError
+from a2a.client.errors import (
+    A2AClientInvalidArgsError,
+    A2AClientInvalidStateError,
+)
 from a2a.server.events.event_queue import Event
 from a2a.types import (
     Message,
@@ -12,7 +15,6 @@
 )
 from a2a.utils import append_artifact_to_task
 
-
 logger = logging.getLogger(__name__)
 
 
@@ -45,6 +47,24 @@ def get_task(self) -> Task | None:
 
         return self._current_task
 
+    def get_task_or_raise(self) -> Task:
+        """Retrieves the current task object.
+
+        Returns:
+            The `Task` object.
+
+        Raises:
+            A2AClientInvalidStateError: If there is no current known Task.
+        """
+        if not (task := self.get_task()):
+            # Note: The source of this error is either from bad client usage
+            # or from the server sending invalid updates. It indicates that this
+            # task manager has not consumed any information about a task, yet
+            # the caller is attempting to retrieve the current state of the task
+            # it expects to be present.
+            raise A2AClientInvalidStateError('no current Task')
+        return task
+
     async def save_task_event(
         self, event: Task | TaskStatusUpdateEvent | TaskArtifactUpdateEvent
     ) -> Task | None:

src/a2a/client/errors.py

@@ -1,5 +1,7 @@
 """Custom exceptions for the A2A client."""
 
+from a2a.types import JSONRPCErrorResponse
+
 
 class A2AClientError(Exception):
     """Base exception for A2A Client errors."""
@@ -57,3 +59,31 @@ def __init__(self, message: str):
         """
         self.message = message
         super().__init__(f'Invalid arguments error: {message}')
+
+
+class A2AClientInvalidStateError(A2AClientError):
+    """Client exception for an invalid client state."""
+
+    def __init__(self, message: str):
+        """Initializes the A2AClientInvalidStateError.
+
+        Args:
+            message: A descriptive error message.
+        """
+        self.message = message
+        super().__init__(f'Invalid state error: {message}')
+
+
+class A2AClientJSONRPCError(A2AClientError):
+    """Client exception for JSON-RPC errors returned by the server."""
+
+    def __init__(self, error: JSONRPCErrorResponse):
+        """Initializes the A2AClientJsonRPCError.
+
+        Args:
+            code: The error code.
+            message: A descriptive error message.
+            data: Optional additional error data.
+        """
+        self.error = error.error
+        super().__init__(f'JSON-RPC Error {error.error}')

src/a2a/client/grpc_client.py

@@ -21,6 +21,7 @@
     Consumer,
 )
 from a2a.client.client_task_manager import ClientTaskManager
+from a2a.client.errors import A2AClientInvalidStateError
 from a2a.client.middleware import ClientCallInterceptor
 from a2a.grpc import a2a_pb2, a2a_pb2_grpc
 from a2a.types import (
@@ -104,8 +105,7 @@ async def send_message_streaming(
         *,
         context: ClientCallContext | None = None,
     ) -> AsyncGenerator[
-        Message | Task | TaskStatusUpdateEvent | TaskArtifactUpdateEvent,
-        None,
+        Message | Task | TaskStatusUpdateEvent | TaskArtifactUpdateEvent
     ]:
         """Sends a streaming message request to the agent and yields responses as they arrive.
 
@@ -134,18 +134,36 @@ async def send_message_streaming(
             response = await stream.read()
             if response == grpc.aio.EOF:  # pyright: ignore [reportAttributeAccessIssue]
                 break
-            if response.HasField('msg'):
-                yield proto_utils.FromProto.message(response.msg)
-            elif response.HasField('task'):
-                yield proto_utils.FromProto.task(response.task)
-            elif response.HasField('status_update'):
-                yield proto_utils.FromProto.task_status_update_event(
-                    response.status_update
-                )
-            elif response.HasField('artifact_update'):
-                yield proto_utils.FromProto.task_artifact_update_event(
-                    response.artifact_update
-                )
+            yield proto_utils.FromProto.stream_response(response)
+
+    async def resubscribe(
+        self, request: TaskIdParams, *, context: ClientCallContext | None = None
+    ) -> AsyncGenerator[
+        Task | Message | TaskStatusUpdateEvent | TaskArtifactUpdateEvent
+    ]:
+        """Reconnects to get task updates.
+
+        This method uses a unary server-side stream to receive updates.
+
+        Args:
+            request: The `TaskIdParams` object containing the task information to reconnect to.
+            context: The client call context.
+
+        Yields:
+            Task update events, which can be either a Task, Message,
+            TaskStatusUpdateEvent, or TaskArtifactUpdateEvent.
+
+        Raises:
+            A2AClientInvalidStateError: If the server returns an invalid response.
+        """
+        stream = self.stub.TaskSubscription(
+            a2a_pb2.TaskSubscriptionRequest(name=f'tasks/{request.id}')
+        )
+        while True:
+            response = await stream.read()
+            if response == grpc.aio.EOF:  # pyright: ignore [reportAttributeAccessIssue]
+                break
+            yield proto_utils.FromProto.stream_response(response)
 
     async def get_task(
         self,
@@ -283,9 +301,7 @@ def __init__(
             raise ValueError('GRPC client requires channel factory.')
         self._card = card
         self._config = config
-        # Defer init to first use.
-        self._transport_client = None
-        channel = self._config.grpc_channel_factory(self._card.url)
+        channel = config.grpc_channel_factory(self._card.url)
         stub = a2a_pb2_grpc.A2AServiceStub(channel)
         self._transport_client = GrpcTransportClient(stub, self._card)
 
@@ -331,27 +347,45 @@ async def send_message(
             await self.consume(result, self._card)
             yield result
             return
-        # Get Task tracker
         tracker = ClientTaskManager()
-        async for event in self._transport_client.send_message_streaming(
+        stream = self._transport_client.send_message_streaming(
             MessageSendParams(
                 message=request,
                 configuration=config,
             ),
             context=context,
-        ):
-            # Update task, check for errors, etc.
-            if isinstance(event, Message):
-                await self.consume(event, self._card)
-                yield event
-                return
-            await tracker.process(event)
-            result = (
-                tracker.get_task(),
-                None if isinstance(event, Task) else event,
+        )
+        # Only the first event may be a Message. All others must be Task
+        # or TaskStatusUpdates. Separate this one out, which allows our core
+        # event processing logic to ignore that case.
+        # TODO(mikeas1): Reconcile with other transport logic.
+        first_event = await anext(stream)
+        if isinstance(first_event, Message):
+            yield first_event
+            return
+        yield await self._process_response(tracker, first_event)
+        async for result in stream:
+            yield await self._process_response(tracker, result)
+
+    async def _process_response(
+        self,
+        tracker: ClientTaskManager,
+        event: Task | Message | TaskStatusUpdateEvent | TaskArtifactUpdateEvent,
+    ) -> ClientEvent:
+        result = event.root.result
+        # Update task, check for errors, etc.
+        if isinstance(result, Message):
+            raise A2AClientInvalidStateError(
+                'received a streamed Message from server after first response; this'
+                ' is not supported'
             )
-            await self.consume(result, self._card)
-            yield result
+        await tracker.process(result)
+        result = (
+            tracker.get_task_or_raise(),
+            None if isinstance(result, Task) else result,
+        )
+        await self.consume(result, self._card)
+        return result
 
     async def get_task(
         self,
@@ -438,7 +472,7 @@ async def resubscribe(
         request: TaskIdParams,
         *,
         context: ClientCallContext | None = None,
-    ) -> AsyncIterator[Task | Message]:
+    ) -> AsyncIterator[ClientEvent]:
         """Resubscribes to a task's event stream.
 
         This is only available if both the client and server support streaming.
@@ -464,12 +498,14 @@ async def resubscribe(
             raise NotImplementedError(
                 'Resubscribe is not implemented on the gRPC transport client.'
             )
-        async for event in self._transport_client.resubscribe(
+        # Note: works correctly for resubscription where the first event is the
+        # current Task state.
+        tracker = ClientTaskManager()
+        async for result in self._transport_client.resubscribe(
             request,
             context=context,
         ):
-            # Update task, check for errors, etc.
-            yield event
+            yield await self._process_response(tracker, result)
 
     async def get_card(
         self,