[DRAFT] feat: Upgrade A2A to v1.0

This updates the SDK to be A2A v1.0 compliant, all types are generated
from the v1.0 a2a.proto. JSONRPC/HTTP+JSON transports are converted to
use the a2a types encoded using ProtoJSON directly from the generated
types.

Author: Tehsmash

buf.gen.yaml

@@ -2,7 +2,7 @@
 version: v2
 inputs:
   - git_repo: https://github.com/a2aproject/A2A.git
-    ref: main
+    ref: transports
     subdir: specification/grpc
 managed:
   enabled: true
@@ -21,11 +21,11 @@ plugins:
   # Generate python protobuf related code
   # Generates *_pb2.py files, one for each .proto
   - remote: buf.build/protocolbuffers/python:v29.3
-    out: src/a2a/grpc
+    out: src/a2a/types
   # Generate python service code.
   # Generates *_pb2_grpc.py
   - remote: buf.build/grpc/python
-    out: src/a2a/grpc
+    out: src/a2a/types
   # Generates *_pb2.pyi files.
   - remote: buf.build/protocolbuffers/pyi
-    out: src/a2a/grpc
+    out: src/a2a/types

pyproject.toml

@@ -114,7 +114,7 @@ explicit = true
 
 [tool.mypy]
 plugins = ["pydantic.mypy"]
-exclude = ["src/a2a/grpc/"]
+exclude = ["src/a2a/types/a2a_pb2\\.py", "src/a2a/types/a2a_pb2_grpc\\.py"]
 disable_error_code = [
   "import-not-found",
   "annotation-unchecked",
@@ -134,7 +134,8 @@ exclude = [
   "**/node_modules",
   "**/venv",
   "**/.venv",
-  "src/a2a/grpc/",
+  "src/a2a/types/a2a_pb2.py",
+  "src/a2a/types/a2a_pb2_grpc.py",
 ]
 reportMissingImports = "none"
 reportMissingModuleSource = "none"
@@ -145,7 +146,8 @@ omit = [
     "*/tests/*",
     "*/site-packages/*",
     "*/__init__.py",
-    "src/a2a/grpc/*",
+    "src/a2a/types/a2a_pb2.py",
+    "src/a2a/types/a2a_pb2_grpc.py",
 ]
 
 [tool.coverage.report]
@@ -257,7 +259,9 @@ exclude = [
     "node_modules",
     "venv",
     "*/migrations/*",
-    "src/a2a/grpc/**",
+    "src/a2a/types/a2a_pb2.py",
+    "src/a2a/types/a2a_pb2.pyi",
+    "src/a2a/types/a2a_pb2_grpc.py",
     "tests/**",
 ]
 
@@ -311,7 +315,9 @@ inline-quotes = "single"
 
 [tool.ruff.format]
 exclude = [
-    "src/a2a/grpc/**",
+    "src/a2a/types/a2a_pb2.py",
+    "src/a2a/types/a2a_pb2.pyi",
+    "src/a2a/types/a2a_pb2_grpc.py",
 ]
 docstring-code-format = true
 docstring-code-line-length = "dynamic"

src/a2a/client/auth/interceptor.py

@@ -3,7 +3,7 @@
 
 from a2a.client.auth.credentials import CredentialService
 from a2a.client.middleware import ClientCallContext, ClientCallInterceptor
-from a2a.types import (
+from a2a.types.a2a_pb2 import (
     AgentCard,
     APIKeySecurityScheme,
     HTTPAuthSecurityScheme,

src/a2a/client/base_client.py

@@ -1,29 +1,31 @@
-from collections.abc import AsyncIterator
+from collections.abc import AsyncIterator, AsyncGenerator
 from typing import Any
 
 from a2a.client.client import (
     Client,
     ClientCallContext,
     ClientConfig,
-    ClientEvent,
     Consumer,
+    ClientEvent,
 )
 from a2a.client.client_task_manager import ClientTaskManager
 from a2a.client.errors import A2AClientInvalidStateError
 from a2a.client.middleware import ClientCallInterceptor
 from a2a.client.transports.base import ClientTransport
-from a2a.types import (
+from a2a.types.a2a_pb2 import (
     AgentCard,
     GetTaskPushNotificationConfigParams,
     Message,
-    MessageSendConfiguration,
-    MessageSendParams,
+    SendMessageConfiguration,
+    SendMessageRequest,
     Task,
     TaskArtifactUpdateEvent,
-    TaskIdParams,
+    SubscribeToTaskRequest,
+    CancelTaskRequest,
     TaskPushNotificationConfig,
-    TaskQueryParams,
+    GetTaskRequest,
     TaskStatusUpdateEvent,
+    StreamResponse,
 )
 
 
@@ -50,7 +52,7 @@ async def send_message(
         context: ClientCallContext | None = None,
         request_metadata: dict[str, Any] | None = None,
         extensions: list[str] | None = None,
-    ) -> AsyncIterator[ClientEvent | Message]:
+    ) -> AsyncIterator[StreamResponse]:
         """Sends a message to the agent.
 
         This method handles both streaming and non-streaming (polling) interactions
@@ -64,9 +66,9 @@ async def send_message(
             extensions: List of extensions to be activated.
 
         Yields:
-            An async iterator of `ClientEvent` or a final `Message` response.
+            An async iterator of `ClientEvent`
         """
-        config = MessageSendConfiguration(
+        config = SendMessageConfiguration(
             accepted_output_modes=self._config.accepted_output_modes,
             blocking=not self._config.polling,
             push_notification_config=(
@@ -75,67 +77,63 @@ async def send_message(
                 else None
             ),
         )
-        params = MessageSendParams(
-            message=request, configuration=config, metadata=request_metadata
+        params = SendMessageRequest(
+            request=request, configuration=config, metadata=request_metadata
         )
 
         if not self._config.streaming or not self._card.capabilities.streaming:
             response = await self._transport.send_message(
                 params, context=context, extensions=extensions
             )
-            result = (
-                (response, None) if isinstance(response, Task) else response
-            )
-            await self.consume(result, self._card)
-            yield result
+
+            # In non-streaming case we convert to a StreamResponse so that the
+            # client always sees the same iterator.
+            stream_response = StreamResponse()
+            if response.HasField("task"):
+                stream_response.task = response.task
+                client_event = (stream_response, response.task)
+
+            if response.HasField("message"):
+                stream_response.message = response.message
+                client_event = (stream_response, None)
+
+            await self.consume(client_event, self._card)
+            yield client_event
             return
 
-        tracker = ClientTaskManager()
         stream = self._transport.send_message_streaming(
             params, context=context, extensions=extensions
         )
+        async for client_event in self._process_stream(stream):
+            yield client_event
 
-        first_event = await anext(stream)
-        # The response from a server may be either exactly one Message or a
-        # series of Task updates. Separate out the first message for special
-        # case handling, which allows us to simplify further stream processing.
-        if isinstance(first_event, Message):
-            await self.consume(first_event, self._card)
-            yield first_event
-            return
-
-        yield await self._process_response(tracker, first_event)
+    async def _process_stream(self, stream: AsyncIterator[StreamResponse]) -> AsyncGenerator[ClientEvent]:
+        tracker = ClientTaskManager()
+        async for stream_response in stream:
+            await self.consume(stream_response)
 
-        async for event in stream:
-            yield await self._process_response(tracker, event)
+            # When we get a message in the stream then we don't expect any
+            # further messages so yield and return
+            if stream_response.HasField("message"):
+                yield stream_response, None
+                return
 
-    async def _process_response(
-        self,
-        tracker: ClientTaskManager,
-        event: Task | Message | TaskStatusUpdateEvent | TaskArtifactUpdateEvent,
-    ) -> ClientEvent:
-        if isinstance(event, Message):
-            raise A2AClientInvalidStateError(
-                'received a streamed Message from server after first response; this is not supported'
-            )
-        await tracker.process(event)
-        task = tracker.get_task_or_raise()
-        update = None if isinstance(event, Task) else event
-        client_event = (task, update)
-        await self.consume(client_event, self._card)
-        return client_event
+            # Otherwise track the task / task update then yield to the client
+            tracker.process(stream_response)
+            updated_task = tracker.get_task_or_raise()
+            yield stream_response, updated_task
 
     async def get_task(
         self,
-        request: TaskQueryParams,
+        request: GetTaskRequest,
         *,
         context: ClientCallContext | None = None,
         extensions: list[str] | None = None,
     ) -> Task:
         """Retrieves the current state and history of a specific task.
 
         Args:
-            request: The `TaskQueryParams` object specifying the task ID.
+            request: The `GetTaskRequest` object specifying the task ID.
             context: The client call context.
             extensions: List of extensions to be activated.
 
@@ -148,15 +146,15 @@ async def get_task(
 
     async def cancel_task(
         self,
-        request: TaskIdParams,
+        request: CancelTaskRequest,
         *,
         context: ClientCallContext | None = None,
         extensions: list[str] | None = None,
     ) -> Task:
         """Requests the agent to cancel a specific task.
 
         Args:
-            request: The `TaskIdParams` object specifying the task ID.
+            request: The `CancelTaskRequest` object specifying the task ID.
             context: The client call context.
             extensions: List of extensions to be activated.
 
@@ -211,7 +209,7 @@ async def get_task_callback(
 
     async def resubscribe(
         self,
-        request: TaskIdParams,
+        request: SubscribeToTaskRequest,
         *,
         context: ClientCallContext | None = None,
         extensions: list[str] | None = None,
@@ -240,10 +238,11 @@ async def resubscribe(
         # Note: resubscribe can only be called on an existing task. As such,
         # we should never see Message updates, despite the typing of the service
         # definition indicating it may be possible.
-        async for event in self._transport.resubscribe(
+        stream = self._transport.resubscribe(
             request, context=context, extensions=extensions
-        ):
-            yield await self._process_response(tracker, event)
+        )
+        async for client_event in self._process_stream(stream):
+            yield client_event
 
     async def get_card(
         self,

src/a2a/client/card_resolver.py

@@ -11,7 +11,7 @@
     A2AClientHTTPError,
     A2AClientJSONError,
 )
-from a2a.types import (
+from a2a.types.a2a_pb2 import (
     AgentCard,
 )
 from a2a.utils.constants import AGENT_CARD_WELL_KNOWN_PATH

src/a2a/client/client.py

@@ -9,7 +9,7 @@
 
 from a2a.client.middleware import ClientCallContext, ClientCallInterceptor
 from a2a.client.optionals import Channel
-from a2a.types import (
+from a2a.types.a2a_pb2 import (
     AgentCard,
     GetTaskPushNotificationConfigParams,
     Message,
@@ -21,6 +21,7 @@
     TaskQueryParams,
     TaskStatusUpdateEvent,
     TransportProtocol,
+    StreamResponse,
 )
 
 
@@ -71,14 +72,11 @@ class ClientConfig:
     """A list of extension URIs the client supports."""
 
 
-UpdateEvent = TaskStatusUpdateEvent | TaskArtifactUpdateEvent | None
-# Alias for emitted events from client
-ClientEvent = tuple[Task, UpdateEvent]
+ClientEvent = [StreamResponse, Task | None]
+
 # Alias for an event consuming callback. It takes either a (task, update) pair
 # or a message as well as the agent card for the agent this came from.
-Consumer = Callable[
-    [ClientEvent | Message, AgentCard], Coroutine[None, Any, Any]
-]
+Consumer = Callable[ClientEvent, Coroutine[None, Any, Any]]
 
 
 class Client(ABC):
@@ -115,7 +113,7 @@ async def send_message(
         context: ClientCallContext | None = None,
         request_metadata: dict[str, Any] | None = None,
         extensions: list[str] | None = None,
-    ) -> AsyncIterator[ClientEvent | Message]:
+    ) -> AsyncIterator[StreamResponse]:
         """Sends a message to the server.
 
         This will automatically use the streaming or non-streaming approach
@@ -174,7 +172,7 @@ async def resubscribe(
         *,
         context: ClientCallContext | None = None,
         extensions: list[str] | None = None,
-    ) -> AsyncIterator[ClientEvent]:
+    ) -> AsyncIterator[StreamResponse]:
         """Resubscribes to a task's event stream."""
         return
         yield