feat: add plugin framework

.github/actions/spelling/allow.txt

@@ -1,4 +1,4 @@
 a2a
 A2A
 A2AFastAPI
 AAgent
@@ -29,7 +29,7 @@
 AUser
 autouse
 backticks
 base64url
 buf
 bufbuild
 cla
@@ -49,7 +49,7 @@
 dsn
 DSNs
 dunders
 ES256
 euo
 EUR
 evt
@@ -66,8 +66,8 @@
 gowebpki
 GVsb
 hazmat
 HS256
 HS384
 ietf
 importlib
 initdb
@@ -110,11 +110,12 @@
 Oneof
 OpenAPI
 openapiv
 openapiv2
 opensource
 otherurl
 outerjoin
 pb2
+PLW
 podman
 Podman
 poolclass
@@ -138,7 +139,7 @@
 respx
 resub
 rmi
 RS256
 RUF
 Rundgren
 SECP

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_client_request`` before
+    the interceptors, ``after_client_request`` 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,
@@ -90,15 +110,20 @@ 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',
+                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
@@ -324,15 +349,18 @@ 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',
+                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,
@@ -372,6 +400,19 @@ async def close(self) -> None:
         """Closes the underlying transport."""
         await self._transport.close()
 
+    def _client_request(
+        self,
+        input_data: Any,
+        method: str,
+        context: ClientCallContext | None,
+    ) -> ClientRequest:
+        return ClientRequest(
+            method=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 +421,41 @@ async def _execute_with_interceptors(
         transport_call: Callable[
             [Any, ClientCallContext | None], Awaitable[Any]
         ],
+    ) -> Any:
+        """Runs a unary call through plugins, interceptors and the transport.
+
+        ``method`` is reported to both plugins and interceptors.
+        """
+        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, method, context)
+        await plugins.before_client_request(request)
+        try:
+            result = await self._call_through_interceptors(
+                request.payload, method, request.context, transport_call
+            )
+        except Exception as e:
+            await plugins.after_client_request(
+                ClientResponse(request.method, None, request.context, error=e)
+            )
+            raise
+        await plugins.after_client_request(
+            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 +498,62 @@ async def _execute_stream_with_interceptors(
         transport_call: Callable[
             [Any, ClientCallContext | None], AsyncIterator[StreamResponse]
         ],
-    ) -> AsyncIterator[StreamResponse]:
+    ) -> AsyncGenerator[StreamResponse, None]:
+        """Runs a streaming call through plugins, interceptors and the transport.
 
+        ``method`` is reported to both plugins and interceptors.
+        """
+        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, method, context)
+        await plugins.before_client_request(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_client_request(
+                        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_client_request(
+                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',
+]