feat: add plugin framework

src/a2a/client/base_client.py

@@ -1,4 +1,5 @@
 from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable
+from contextlib import aclosing
 from typing import Any
 
 from a2a.client.client import (
@@ -12,6 +13,8 @@
     ClientCallInterceptor,
 )
 from a2a.client.transports.base import ClientTransport
+from a2a.plugins.manager import PluginManager, get_global
+from a2a.plugins.types import ClientRequest, ClientResponse
 from a2a.types.a2a_pb2 import (
     AgentCard,
     CancelTaskRequest,
@@ -32,20 +35,37 @@
 
 
 class BaseClient(Client):
-    """Base implementation of the A2A client, containing transport-independent logic."""
+    """Base implementation of the A2A client, containing transport-independent logic.
+
+    Client plugin hooks run around every call: ``before_agent_call`` before
+    the interceptors, ``after_agent_call`` after them.
+    """
 
     def __init__(
         self,
         card: AgentCard,
         config: ClientConfig,
         transport: ClientTransport,
         interceptors: list[ClientCallInterceptor],
+        *,
+        plugins: PluginManager | None = None,
     ):
+        """Initializes the client.
+
+        Args:
+            card: The card of the agent to call.
+            config: The client configuration.
+            transport: The transport used for calls.
+            interceptors: Interceptors applied to every call.
+            plugins: Plugin manager whose client hooks run for every call.
+                Defaults to the global manager at construction time.
+        """
         super().__init__(interceptors)
         self._card = card
         self._config = config
         self._transport = transport
         self._interceptors = interceptors
+        self._plugins = plugins if plugins is not None else get_global()
 
     async def send_message(
         self,
@@ -71,6 +91,7 @@ async def send_message(
             response = await self._execute_with_interceptors(
                 input_data=request,
                 method='send_message',
+                rpc_method='SendMessage',
                 context=context,
                 transport_call=lambda req, ctx: self._transport.send_message(
                     req, context=ctx
@@ -90,15 +111,21 @@ async def send_message(
             yield stream_response
             return
 
-        async for event in self._execute_stream_with_interceptors(
-            input_data=request,
-            method='send_message_streaming',
-            context=context,
-            transport_call=lambda req, ctx: (
-                self._transport.send_message_streaming(req, context=ctx)
-            ),
-        ):
-            yield event
+        # aclosing() closes the inner stream (and the transport stream) as
+        # soon as the caller stops iterating.
+        async with aclosing(
+            self._execute_stream_with_interceptors(
+                input_data=request,
+                method='send_message_streaming',
+                rpc_method='SendStreamingMessage',
+                context=context,
+                transport_call=lambda req, ctx: (
+                    self._transport.send_message_streaming(req, context=ctx)
+                ),
+            )
+        ) as events:
+            async for event in events:
+                yield event
 
     def _apply_client_config(self, request: SendMessageRequest) -> None:
         request.configuration.return_immediately |= self._config.polling
@@ -152,6 +179,7 @@ async def get_task(
         return await self._execute_with_interceptors(
             input_data=request,
             method='get_task',
+            rpc_method='GetTask',
             context=context,
             transport_call=lambda req, ctx: self._transport.get_task(
                 req, context=ctx
@@ -168,6 +196,7 @@ async def list_tasks(
         return await self._execute_with_interceptors(
             input_data=request,
             method='list_tasks',
+            rpc_method='ListTasks',
             context=context,
             transport_call=lambda req, ctx: self._transport.list_tasks(
                 req, context=ctx
@@ -192,6 +221,7 @@ async def cancel_task(
         return await self._execute_with_interceptors(
             input_data=request,
             method='cancel_task',
+            rpc_method='CancelTask',
             context=context,
             transport_call=lambda req, ctx: self._transport.cancel_task(
                 req, context=ctx
@@ -216,6 +246,7 @@ async def create_task_push_notification_config(
         return await self._execute_with_interceptors(
             input_data=request,
             method='create_task_push_notification_config',
+            rpc_method='CreateTaskPushNotificationConfig',
             context=context,
             transport_call=lambda req, ctx: (
                 self._transport.create_task_push_notification_config(
@@ -242,6 +273,7 @@ async def get_task_push_notification_config(
         return await self._execute_with_interceptors(
             input_data=request,
             method='get_task_push_notification_config',
+            rpc_method='GetTaskPushNotificationConfig',
             context=context,
             transport_call=lambda req, ctx: (
                 self._transport.get_task_push_notification_config(
@@ -268,6 +300,7 @@ async def list_task_push_notification_configs(
         return await self._execute_with_interceptors(
             input_data=request,
             method='list_task_push_notification_configs',
+            rpc_method='ListTaskPushNotificationConfigs',
             context=context,
             transport_call=lambda req, ctx: (
                 self._transport.list_task_push_notification_configs(
@@ -291,6 +324,7 @@ async def delete_task_push_notification_config(
         return await self._execute_with_interceptors(
             input_data=request,
             method='delete_task_push_notification_config',
+            rpc_method='DeleteTaskPushNotificationConfig',
             context=context,
             transport_call=lambda req, ctx: (
                 self._transport.delete_task_push_notification_config(
@@ -324,15 +358,19 @@ async def subscribe(
                 'client and/or server do not support resubscription.'
             )
 
-        async for event in self._execute_stream_with_interceptors(
-            input_data=request,
-            method='subscribe',
-            context=context,
-            transport_call=lambda req, ctx: self._transport.subscribe(
-                req, context=ctx
-            ),
-        ):
-            yield event
+        async with aclosing(
+            self._execute_stream_with_interceptors(
+                input_data=request,
+                method='subscribe',
+                rpc_method='SubscribeToTask',
+                context=context,
+                transport_call=lambda req, ctx: self._transport.subscribe(
+                    req, context=ctx
+                ),
+            )
+        ) as events:
+            async for event in events:
+                yield event
 
     async def get_extended_agent_card(
         self,
@@ -357,6 +395,7 @@ async def get_extended_agent_card(
         card = await self._execute_with_interceptors(
             input_data=request,
             method='get_extended_agent_card',
+            rpc_method='GetExtendedAgentCard',
             context=context,
             transport_call=lambda req, ctx: (
                 self._transport.get_extended_agent_card(req, context=ctx)
@@ -372,6 +411,19 @@ async def close(self) -> None:
         """Closes the underlying transport."""
         await self._transport.close()
 
+    def _client_request(
+        self,
+        input_data: Any,
+        rpc_method: str,
+        context: ClientCallContext | None,
+    ) -> ClientRequest:
+        return ClientRequest(
+            method=rpc_method,
+            payload=input_data,
+            agent_card=self._card,
+            context=context if context is not None else ClientCallContext(),
+        )
+
     async def _execute_with_interceptors(
         self,
         input_data: Any,
@@ -380,6 +432,46 @@ async def _execute_with_interceptors(
         transport_call: Callable[
             [Any, ClientCallContext | None], Awaitable[Any]
         ],
+        *,
+        rpc_method: str | None = None,
+    ) -> Any:
+        """Runs a unary call through plugins, interceptors and the transport.
+
+        ``method`` is passed to interceptors; ``rpc_method`` is the A2A RPC
+        name reported to plugins and defaults to ``method``.
+        """
+        plugins = self._plugins
+        if not plugins.has_client_hooks:
+            return await self._call_through_interceptors(
+                input_data, method, context, transport_call
+            )
+
+        request = self._client_request(
+            input_data, rpc_method or method, context
+        )
+        await plugins.before_agent_call(request)
+        try:
+            result = await self._call_through_interceptors(
+                request.payload, method, request.context, transport_call
+            )
+        except Exception as e:
+            await plugins.after_agent_call(
+                ClientResponse(request.method, None, request.context, error=e)
+            )
+            raise
+        await plugins.after_agent_call(
+            ClientResponse(request.method, result, request.context)
+        )
+        return result
+
+    async def _call_through_interceptors(
+        self,
+        input_data: Any,
+        method: str,
+        context: ClientCallContext | None,
+        transport_call: Callable[
+            [Any, ClientCallContext | None], Awaitable[Any]
+        ],
     ) -> Any:
         before_args = BeforeArgs(
             input=input_data,
@@ -422,8 +514,67 @@ async def _execute_stream_with_interceptors(
         transport_call: Callable[
             [Any, ClientCallContext | None], AsyncIterator[StreamResponse]
         ],
-    ) -> AsyncIterator[StreamResponse]:
+        *,
+        rpc_method: str | None = None,
+    ) -> AsyncGenerator[StreamResponse, None]:
+        """Runs a streaming call through plugins, interceptors and the transport.
 
+        ``method`` is passed to interceptors; ``rpc_method`` is the A2A RPC
+        name reported to plugins and defaults to ``method``.
+        """
+        plugins = self._plugins
+        if not plugins.has_client_hooks:
+            async with aclosing(
+                self._stream_through_interceptors(
+                    input_data, method, context, transport_call
+                )
+            ) as events:
+                async for event in events:
+                    yield event
+            return
+
+        request = self._client_request(
+            input_data, rpc_method or method, context
+        )
+        await plugins.before_agent_call(request)
+        error: BaseException | None = None
+        try:
+            async with aclosing(
+                self._stream_through_interceptors(
+                    request.payload, method, request.context, transport_call
+                )
+            ) as events:
+                async for event in events:
+                    await plugins.after_agent_call(
+                        ClientResponse(
+                            request.method,
+                            event,
+                            request.context,
+                            is_stream_event=True,
+                            is_final=False,
+                        )
+                    )
+                    yield event
+        except Exception as e:
+            error = e
+            raise
+        finally:
+            # Also runs when the caller stops iterating early.
+            await plugins.after_agent_call(
+                ClientResponse(
+                    request.method, None, request.context, error=error
+                )
+            )
+
+    async def _stream_through_interceptors(
+        self,
+        input_data: Any,
+        method: str,
+        context: ClientCallContext | None,
+        transport_call: Callable[
+            [Any, ClientCallContext | None], AsyncIterator[StreamResponse]
+        ],
+    ) -> AsyncGenerator[StreamResponse, None]:
         before_args = BeforeArgs(
             input=input_data,
             method=method,

src/a2a/plugins/__init__.py

@@ -0,0 +1,63 @@
+"""Plugins: hooks into the A2A client and `DefaultRequestHandlerV2`.
+
+Example:
+    from a2a import plugins
+
+    class TaskAudit(plugins.TaskObserverPlugin):
+        name = 'task-audit'
+
+        async def on_task_event(self, info: plugins.TaskEventInfo) -> None:
+            print(info.task_id, info.from_state, '->', info.to_state)
+
+    plugins.set_global(plugins.sequential(TaskAudit()))
+    # Create DefaultRequestHandlerV2 and clients after this point.
+"""
+
+from a2a.plugins.base import (
+    ClientPlugin,
+    CompletePlugin,
+    ExecutionPlugin,
+    Plugin,
+    PushPlugin,
+    ServerPlugin,
+    TaskObserverPlugin,
+)
+from a2a.plugins.manager import (
+    PluginManager,
+    SequentialPluginManager,
+    get_global,
+    sequential,
+    set_global,
+)
+from a2a.plugins.types import (
+    ClientRequest,
+    ClientResponse,
+    PushRequest,
+    PushResult,
+    ServerRequest,
+    ServerResponse,
+    TaskEventInfo,
+)
+
+
+__all__ = [
+    'ClientPlugin',
+    'ClientRequest',
+    'ClientResponse',
+    'CompletePlugin',
+    'ExecutionPlugin',
+    'Plugin',
+    'PluginManager',
+    'PushPlugin',
+    'PushRequest',
+    'PushResult',
+    'SequentialPluginManager',
+    'ServerPlugin',
+    'ServerRequest',
+    'ServerResponse',
+    'TaskEventInfo',
+    'TaskObserverPlugin',
+    'get_global',
+    'sequential',
+    'set_global',
+]