From ac25825012c20772f95c61a4b3d39fb50d7df648 Mon Sep 17 00:00:00 2001 From: Mykyta Netipa Date: Sat, 17 Oct 2026 15:17:39 +0000 Subject: [PATCH 01/11] feat(plugins): add plugin interfaces, hook types and SequentialPluginManager Implements the plugin model from go/a2a-plugins for the Python SDK: a plugin is one object that implements any of ServerPlugin, ClientPlugin, ExecutionPlugin, TaskObserverPlugin and PushPlugin, so cross-cutting concerns such as telemetry or auth plug in with one registration call instead of per-SDK interceptor wiring. SequentialPluginManager runs before-hooks in registration order and after-hooks in reverse, unwinds plugins that already ran when a before-hook rejects a call, and logs and swallows after-hook and observer failures so a buggy plugin cannot break a request. Its has_* properties let call sites skip all plugin work when nothing is registered. --- src/a2a/plugins/__init__.py | 63 ++++++ src/a2a/plugins/_runtime.py | 40 ++++ src/a2a/plugins/base.py | 104 ++++++++++ src/a2a/plugins/manager.py | 291 ++++++++++++++++++++++++++++ src/a2a/plugins/types.py | 177 +++++++++++++++++ tests/plugins/test_manager.py | 347 ++++++++++++++++++++++++++++++++++ 6 files changed, 1022 insertions(+) create mode 100644 src/a2a/plugins/__init__.py create mode 100644 src/a2a/plugins/_runtime.py create mode 100644 src/a2a/plugins/base.py create mode 100644 src/a2a/plugins/manager.py create mode 100644 src/a2a/plugins/types.py create mode 100644 tests/plugins/test_manager.py diff --git a/src/a2a/plugins/__init__.py b/src/a2a/plugins/__init__.py new file mode 100644 index 000000000..e233773a3 --- /dev/null +++ b/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', +] diff --git a/src/a2a/plugins/_runtime.py b/src/a2a/plugins/_runtime.py new file mode 100644 index 000000000..b3fe4378d --- /dev/null +++ b/src/a2a/plugins/_runtime.py @@ -0,0 +1,40 @@ +"""Internal plumbing for scoping plugin hooks to `DefaultRequestHandlerV2`.""" + +from __future__ import annotations + +import contextlib + +from contextvars import ContextVar +from typing import TYPE_CHECKING + + +if TYPE_CHECKING: + from collections.abc import Iterator + + from a2a.plugins.manager import PluginManager + + +# Set by `ActiveTask` around push delivery, so `BasePushNotificationSender` +# runs push hooks for the V2 path only. Child tasks created inside the scope +# (e.g. by asyncio.gather) inherit the value. +_push_plugins: ContextVar[PluginManager | None] = ContextVar( + 'a2a_push_plugins', default=None +) + + +@contextlib.contextmanager +def push_plugins_scope(manager: PluginManager | None) -> Iterator[None]: + """Makes ``manager`` the push plugin manager inside the block.""" + token = _push_plugins.set(manager) + try: + yield + finally: + _push_plugins.reset(token) + + +def current_push_plugins() -> PluginManager | None: + """Returns the manager set by `push_plugins_scope`, if it has push hooks.""" + manager = _push_plugins.get() + if manager is None or not manager.has_push_hooks: + return None + return manager diff --git a/src/a2a/plugins/base.py b/src/a2a/plugins/base.py new file mode 100644 index 000000000..4e3e391ae --- /dev/null +++ b/src/a2a/plugins/base.py @@ -0,0 +1,104 @@ +"""Plugin interfaces. + +A plugin inherits from the interfaces it needs and overrides only the hooks +it uses; every base method is a no-op. + +Interceptor hooks (``before_*``) may modify their argument or reject the call +by raising. Exceptions from ``after_*`` and observer hooks (``on_*``) are +logged and swallowed by the plugin manager. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + + +if TYPE_CHECKING: + from a2a.plugins.types import ( + ClientRequest, + ClientResponse, + PushRequest, + PushResult, + ServerRequest, + ServerResponse, + TaskEventInfo, + ) + from a2a.server.agent_execution.context import RequestContext + from a2a.server.events.event_queue import Event + + +class Plugin: + """Common base of all plugins.""" + + name: str = 'unnamed' + """Name used in logs.""" + + +class ServerPlugin(Plugin): + """Hooks around each request handled by `DefaultRequestHandlerV2`.""" + + async def before_agent_run(self, request: ServerRequest) -> None: + """Called before the handler processes a request. + + Raise to reject the request; the exception is returned to the caller. + """ + + async def after_agent_run(self, response: ServerResponse) -> None: + """Called with the handler result, per stream event, and at the end.""" + + +class ClientPlugin(Plugin): + """Hooks around each call made by an A2A `Client`.""" + + async def before_agent_call(self, request: ClientRequest) -> None: + """Called before the call is sent. Raise to cancel it.""" + + async def after_agent_call(self, response: ClientResponse) -> None: + """Called with the result, per stream event, and at the end.""" + + +class ExecutionPlugin(Plugin): + """Hooks around `AgentExecutor.execute`, which runs outside the request.""" + + async def before_execution(self, request_context: RequestContext) -> None: + """Called in the execution task before the agent runs. + + Raise to fail the execution without running the agent. + """ + + async def on_execution_event( + self, request_context: RequestContext, event: Event + ) -> None: + """Called when the agent enqueues an event, before it is processed. + + Observer: runs on the agent's path, so keep it cheap. + """ + + async def after_execution( + self, request_context: RequestContext, error: BaseException | None + ) -> None: + """Called in the execution task after the agent returns or raises.""" + + +class TaskObserverPlugin(Plugin): + """Observer of task writes.""" + + async def on_task_event(self, info: TaskEventInfo) -> None: + """Called after every successful task write.""" + + +class PushPlugin(Plugin): + """Hooks around each push notification delivery.""" + + async def before_push_send(self, request: PushRequest) -> None: + """Called before the webhook request. + + May add headers or set ``request.skip``. Raising fails the delivery. + """ + + async def on_push_sent(self, result: PushResult) -> None: + """Called after the delivery attempt, or when it was skipped.""" + + +class CompletePlugin(ClientPlugin, ServerPlugin): + """A plugin with both client and server hooks.""" diff --git a/src/a2a/plugins/manager.py b/src/a2a/plugins/manager.py new file mode 100644 index 000000000..c115641da --- /dev/null +++ b/src/a2a/plugins/manager.py @@ -0,0 +1,291 @@ +"""Plugin managers and the global registry.""" + +from __future__ import annotations + +import logging + +from typing import TYPE_CHECKING, TypeVar + +from a2a.plugins.base import ( + ClientPlugin, + ExecutionPlugin, + Plugin, + PushPlugin, + ServerPlugin, + TaskObserverPlugin, +) +from a2a.plugins.types import ClientResponse, ServerResponse + + +if TYPE_CHECKING: + from collections.abc import Awaitable + + from a2a.plugins.types import ( + ClientRequest, + PushRequest, + PushResult, + ServerRequest, + TaskEventInfo, + ) + from a2a.server.agent_execution.context import RequestContext + from a2a.server.events.event_queue import Event + + +logger = logging.getLogger(__name__) + + +async def _run_safely(plugin: Plugin, hook: str, call: Awaitable[None]) -> None: + """Awaits a hook call, logging and swallowing any exception.""" + try: + await call + except Exception: + logger.exception('A2A plugin %r failed in %s', plugin.name, hook) + + +class PluginManager( + ServerPlugin, ClientPlugin, ExecutionPlugin, TaskObserverPlugin, PushPlugin +): + """Runs registered plugins. SDK call sites depend on this type. + + Since a manager implements every plugin interface, managers can be nested. + The ``has_*`` properties let call sites skip all plugin work when no + plugin implements the hooks they would call. + """ + + name = 'plugin-manager' + + @property + def has_server_hooks(self) -> bool: + """Whether any plugin implements `ServerPlugin`.""" + return False + + @property + def has_client_hooks(self) -> bool: + """Whether any plugin implements `ClientPlugin`.""" + return False + + @property + def has_execution_hooks(self) -> bool: + """Whether any plugin implements `ExecutionPlugin`.""" + return False + + @property + def has_task_hooks(self) -> bool: + """Whether any plugin implements `TaskObserverPlugin`.""" + return False + + @property + def has_push_hooks(self) -> bool: + """Whether any plugin implements `PushPlugin`.""" + return False + + +_P = TypeVar('_P', bound=Plugin) + + +def _select( + plugins: tuple[Plugin, ...], kind: type[_P], has_hooks: str +) -> list[_P]: + """Returns plugins implementing ``kind``; nested managers only if non-empty.""" + return [ + p + for p in plugins + if isinstance(p, kind) + and (not isinstance(p, PluginManager) or getattr(p, has_hooks)) + ] + + +class SequentialPluginManager(PluginManager): + """Runs plugins one after another in registration order. + + - ``before_*`` hooks run in registration order. If one raises, plugins + that already ran get their ``after_*`` hook with the error, and the + exception propagates. + - ``after_*`` hooks run in reverse order. Their exceptions are logged and + swallowed. + - Observer hooks (``on_*``) run in registration order. Their exceptions + are logged and swallowed. + """ + + name = 'sequential' + + def __init__(self, *plugins: Plugin): + unknown = [p for p in plugins if not isinstance(p, Plugin)] + if unknown: + raise TypeError(f'Not A2A plugins: {unknown!r}') + self._plugins = plugins + self._server: list[ServerPlugin] = _select( + plugins, ServerPlugin, 'has_server_hooks' + ) + self._client: list[ClientPlugin] = _select( + plugins, ClientPlugin, 'has_client_hooks' + ) + self._execution: list[ExecutionPlugin] = _select( + plugins, ExecutionPlugin, 'has_execution_hooks' + ) + self._task: list[TaskObserverPlugin] = _select( + plugins, TaskObserverPlugin, 'has_task_hooks' + ) + self._push: list[PushPlugin] = _select( + plugins, PushPlugin, 'has_push_hooks' + ) + + @property + def plugins(self) -> tuple[Plugin, ...]: + """The registered plugins, in registration order.""" + return self._plugins + + @property + def has_server_hooks(self) -> bool: # noqa: D102 + return bool(self._server) + + @property + def has_client_hooks(self) -> bool: # noqa: D102 + return bool(self._client) + + @property + def has_execution_hooks(self) -> bool: # noqa: D102 + return bool(self._execution) + + @property + def has_task_hooks(self) -> bool: # noqa: D102 + return bool(self._task) + + @property + def has_push_hooks(self) -> bool: # noqa: D102 + return bool(self._push) + + # --- server --- + + async def before_agent_run(self, request: ServerRequest) -> None: # noqa: D102 + ran: list[ServerPlugin] = [] + for plugin in self._server: + try: + await plugin.before_agent_run(request) + except Exception as e: + response = ServerResponse( + request.method, None, request.context, error=e + ) + for done in reversed(ran): + await _run_safely( + done, 'after_agent_run', done.after_agent_run(response) + ) + raise + ran.append(plugin) + + async def after_agent_run(self, response: ServerResponse) -> None: # noqa: D102 + for plugin in reversed(self._server): + await _run_safely( + plugin, 'after_agent_run', plugin.after_agent_run(response) + ) + + # --- client --- + + async def before_agent_call(self, request: ClientRequest) -> None: # noqa: D102 + ran: list[ClientPlugin] = [] + for plugin in self._client: + try: + await plugin.before_agent_call(request) + except Exception as e: + response = ClientResponse( + request.method, None, request.context, error=e + ) + for done in reversed(ran): + await _run_safely( + done, + 'after_agent_call', + done.after_agent_call(response), + ) + raise + ran.append(plugin) + + async def after_agent_call(self, response: ClientResponse) -> None: # noqa: D102 + for plugin in reversed(self._client): + await _run_safely( + plugin, 'after_agent_call', plugin.after_agent_call(response) + ) + + # --- execution --- + + async def before_execution(self, request_context: RequestContext) -> None: # noqa: D102 + ran: list[ExecutionPlugin] = [] + for plugin in self._execution: + try: + await plugin.before_execution(request_context) + except Exception as e: + for done in reversed(ran): + await _run_safely( + done, + 'after_execution', + done.after_execution(request_context, e), + ) + raise + ran.append(plugin) + + async def on_execution_event( # noqa: D102 + self, request_context: RequestContext, event: Event + ) -> None: + for plugin in self._execution: + await _run_safely( + plugin, + 'on_execution_event', + plugin.on_execution_event(request_context, event), + ) + + async def after_execution( # noqa: D102 + self, request_context: RequestContext, error: BaseException | None + ) -> None: + for plugin in reversed(self._execution): + await _run_safely( + plugin, + 'after_execution', + plugin.after_execution(request_context, error), + ) + + # --- task --- + + async def on_task_event(self, info: TaskEventInfo) -> None: # noqa: D102 + for plugin in self._task: + await _run_safely( + plugin, 'on_task_event', plugin.on_task_event(info) + ) + + # --- push --- + + async def before_push_send(self, request: PushRequest) -> None: # noqa: D102 + for plugin in self._push: + await plugin.before_push_send(request) + + async def on_push_sent(self, result: PushResult) -> None: # noqa: D102 + for plugin in self._push: + await _run_safely( + plugin, 'on_push_sent', plugin.on_push_sent(result) + ) + + +def sequential(*plugins: Plugin) -> SequentialPluginManager: + """Returns a manager that runs ``plugins`` in the given order.""" + return SequentialPluginManager(*plugins) + + +_global_manager: PluginManager = SequentialPluginManager() + + +def set_global(manager: PluginManager) -> None: + """Sets the manager used by handlers and clients created afterwards. + + Handlers and clients capture the manager when they are constructed, so + call this once at startup, before creating them. + """ + global _global_manager # noqa: PLW0603 + if not isinstance(manager, PluginManager): + raise TypeError( + f'Expected a PluginManager, got {type(manager).__name__}; ' + 'wrap plugins with a2a.plugins.sequential(...)' + ) + _global_manager = manager + + +def get_global() -> PluginManager: + """Returns the global manager. Empty unless `set_global` was called.""" + return _global_manager diff --git a/src/a2a/plugins/types.py b/src/a2a/plugins/types.py new file mode 100644 index 000000000..22d48a32e --- /dev/null +++ b/src/a2a/plugins/types.py @@ -0,0 +1,177 @@ +"""Argument types passed to plugin hooks. + +Hooks return ``None``. Per-call state lives in ``context.state`` (server and +client) or ``PushRequest.state`` (push), which the SDK passes unchanged from a +before-hook to the matching after-hook. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any + + +if TYPE_CHECKING: + from a2a.client.client import ClientCallContext + from a2a.server.context import ServerCallContext + from a2a.server.events.event_queue import Event + from a2a.server.tasks.push_notification_sender import ( + PushNotificationEvent, + ) + from a2a.types.a2a_pb2 import AgentCard, Task, TaskState + + +@dataclass +class ServerRequest: + """An incoming request, passed to `ServerPlugin.before_agent_run`. + + Attributes: + method: A2A RPC name, e.g. ``'SendMessage'`` or ``'GetTask'``. + payload: The request proto, e.g. ``SendMessageRequest``. + context: The call context. ``state['headers']`` holds the transport + headers and ``state['a2a.transport']`` the transport name. + agent_card: The agent card served by this handler. + """ + + method: str + payload: Any + context: ServerCallContext + agent_card: AgentCard | None = None + + +@dataclass +class ServerResponse: + """A handler result, passed to `ServerPlugin.after_agent_run`. + + For a unary call the hook runs once with ``is_final=True``. For a stream + it runs once per event with ``is_stream_event=True, is_final=False`` and + once more with ``is_final=True`` and ``payload=None`` when the stream + ends, including on error or client disconnect. + + Attributes: + method: A2A RPC name. + payload: The handler result or stream event, ``None`` on error or + stream end. + context: The same call context passed to `before_agent_run`. + error: The exception raised by the handler, if any. + is_stream_event: Whether ``payload`` is one event of a stream. + is_final: Whether this is the last call for the request. + """ + + method: str + payload: Any + context: ServerCallContext + error: BaseException | None = None + is_stream_event: bool = False + is_final: bool = True + + +@dataclass +class ClientRequest: + """An outgoing call, passed to `ClientPlugin.before_agent_call`. + + Attributes: + method: A2A RPC name, matching the server side. + payload: The request proto. Plugins may replace it. + agent_card: The card of the remote agent. + context: The call context. ``service_parameters`` holds the outgoing + headers. + transport: Protocol binding used for the call, e.g. ``'JSONRPC'``. + url: URL of the selected agent interface. + """ + + method: str + payload: Any + agent_card: AgentCard + context: ClientCallContext + transport: str = '' + url: str = '' + + +@dataclass +class ClientResponse: + """A call result, passed to `ClientPlugin.after_agent_call`. + + Follows the same ``is_stream_event`` / ``is_final`` rules as + `ServerResponse`. + + Attributes: + method: A2A RPC name. + payload: The response or stream event, ``None`` on error or stream end. + context: The same call context passed to `before_agent_call`. + error: The exception raised by the call, if any. + is_stream_event: Whether ``payload`` is one event of a stream. + is_final: Whether this is the last call for the request. + """ + + method: str + payload: Any + context: ClientCallContext + error: BaseException | None = None + is_stream_event: bool = False + is_final: bool = True + + +@dataclass +class TaskEventInfo: + """A successful task write, passed to `TaskObserverPlugin.on_task_event`. + + Attributes: + task_id: ID of the task. + from_state: Persisted state before the write, ``None`` when the write + created the task. + to_state: State after the write. Equal to ``from_state`` for writes + that do not change the state, such as artifact updates. + event: The event that caused the write, ``None`` when the SDK wrote + the task without one. + task: The task after the write. This is the SDK's live object; + copy it if you keep it beyond the hook call. + context: Call context of the request that produced the write. + """ + + task_id: str + from_state: TaskState | None + to_state: TaskState + event: Event | None + task: Task + context: ServerCallContext + + +@dataclass +class PushRequest: + """One push delivery, passed to `PushPlugin.before_push_send`. + + Attributes: + task_id: ID of the task. + url: Webhook URL. + headers: Outgoing HTTP headers. Plugins may add or change entries. + event: The event being pushed. + skip: Set to ``True`` to drop this delivery. + state: Per-delivery plugin state, available again in `PushResult`. + """ + + task_id: str + url: str + headers: dict[str, str] + event: PushNotificationEvent + skip: bool = False + state: dict[str, Any] = field(default_factory=dict) + + +@dataclass +class PushResult: + """The outcome of one push delivery, passed to `PushPlugin.on_push_sent`. + + Attributes: + request: The `PushRequest` for this delivery. + duration_s: Duration of the HTTP request in seconds. + status_code: HTTP status code, ``None`` if no response was received. + error: The exception if delivery failed, ``None`` on success. + skipped: Whether a plugin dropped the delivery. + """ + + request: PushRequest + duration_s: float + status_code: int | None + error: BaseException | None + skipped: bool = False diff --git a/tests/plugins/test_manager.py b/tests/plugins/test_manager.py new file mode 100644 index 000000000..d94c3a9ba --- /dev/null +++ b/tests/plugins/test_manager.py @@ -0,0 +1,347 @@ +from unittest.mock import MagicMock + +import pytest + +from a2a import plugins +from a2a.client.client import ClientCallContext +from a2a.plugins import ( + ClientPlugin, + ClientRequest, + ClientResponse, + CompletePlugin, + ExecutionPlugin, + Plugin, + PushPlugin, + PushRequest, + PushResult, + SequentialPluginManager, + ServerPlugin, + ServerRequest, + ServerResponse, + TaskEventInfo, + TaskObserverPlugin, +) +from a2a.server.context import ServerCallContext +from a2a.types.a2a_pb2 import AgentCard, Task, TaskState + + +class Recorder(CompletePlugin, ExecutionPlugin, TaskObserverPlugin, PushPlugin): + """Appends '.' to a shared log for every hook call.""" + + def __init__(self, name: str, log: list[str], fail: set[str] = frozenset()): + self.name = name + self.log = log + self.fail = fail + + def _record(self, hook: str) -> None: + self.log.append(f'{self.name}.{hook}') + if hook in self.fail: + raise ValueError(f'{self.name}.{hook} failed') + + async def before_agent_run(self, request): + self._record('before_agent_run') + + async def after_agent_run(self, response): + self._record('after_agent_run') + + async def before_agent_call(self, request): + self._record('before_agent_call') + + async def after_agent_call(self, response): + self._record('after_agent_call') + + async def before_execution(self, request_context): + self._record('before_execution') + + async def on_execution_event(self, request_context, event): + self._record('on_execution_event') + + async def after_execution(self, request_context, error): + self._record('after_execution') + + async def on_task_event(self, info): + self._record('on_task_event') + + async def before_push_send(self, request): + self._record('before_push_send') + + async def on_push_sent(self, result): + self._record('on_push_sent') + + +def server_request() -> ServerRequest: + return ServerRequest('GetTask', None, ServerCallContext()) + + +def client_request() -> ClientRequest: + return ClientRequest('GetTask', None, AgentCard(), ClientCallContext()) + + +def push_request() -> PushRequest: + return PushRequest('t-1', 'http://hook', {}, Task(id='t-1')) + + +@pytest.fixture(autouse=True) +def restore_global(): + previous = plugins.get_global() + yield + plugins.set_global(previous) + + +@pytest.mark.asyncio +async def test_before_in_order_after_in_reverse(): + log: list[str] = [] + m = plugins.sequential(Recorder('a', log), Recorder('b', log)) + + await m.before_agent_run(server_request()) + await m.after_agent_run( + ServerResponse('GetTask', None, ServerCallContext()) + ) + await m.before_agent_call(client_request()) + await m.after_agent_call( + ClientResponse('GetTask', None, ClientCallContext()) + ) + await m.before_execution(MagicMock()) + await m.after_execution(MagicMock(), None) + + assert log == [ + 'a.before_agent_run', + 'b.before_agent_run', + 'b.after_agent_run', + 'a.after_agent_run', + 'a.before_agent_call', + 'b.before_agent_call', + 'b.after_agent_call', + 'a.after_agent_call', + 'a.before_execution', + 'b.before_execution', + 'b.after_execution', + 'a.after_execution', + ] + + +@pytest.mark.asyncio +async def test_observers_run_in_order(): + log: list[str] = [] + m = plugins.sequential(Recorder('a', log), Recorder('b', log)) + info = TaskEventInfo( + 't-1', + None, + TaskState.TASK_STATE_SUBMITTED, + None, + Task(id='t-1'), + ServerCallContext(), + ) + + await m.on_task_event(info) + await m.on_execution_event(MagicMock(), Task()) + await m.before_push_send(push_request()) + await m.on_push_sent(PushResult(push_request(), 0.1, 200, None)) + + assert log == [ + 'a.on_task_event', + 'b.on_task_event', + 'a.on_execution_event', + 'b.on_execution_event', + 'a.before_push_send', + 'b.before_push_send', + 'a.on_push_sent', + 'b.on_push_sent', + ] + + +@pytest.mark.parametrize( + ('before', 'after', 'call'), + [ + ( + 'before_agent_run', + 'after_agent_run', + lambda m: m.before_agent_run(server_request()), + ), + ( + 'before_agent_call', + 'after_agent_call', + lambda m: m.before_agent_call(client_request()), + ), + ( + 'before_execution', + 'after_execution', + lambda m: m.before_execution(MagicMock()), + ), + ], +) +@pytest.mark.asyncio +async def test_rejection_unwinds_plugins_that_ran(before, after, call): + log: list[str] = [] + m = plugins.sequential( + Recorder('a', log), + Recorder('b', log, fail={before}), + Recorder('c', log), + ) + + with pytest.raises(ValueError, match=f'b.{before} failed'): + await call(m) + + # c never ran; a ran and is unwound; b raised and gets no after-hook. + assert log == [f'a.{before}', f'b.{before}', f'a.{after}'] + + +@pytest.mark.asyncio +async def test_unwind_passes_error_to_after_hook(): + seen: list[BaseException | None] = [] + + class Capture(ServerPlugin): + async def after_agent_run(self, response): + seen.append(response.error) + + class Reject(ServerPlugin): + async def before_agent_run(self, request): + raise PermissionError('denied') + + m = plugins.sequential(Capture(), Reject()) + with pytest.raises(PermissionError): + await m.before_agent_run(server_request()) + + assert len(seen) == 1 + assert isinstance(seen[0], PermissionError) + + +@pytest.mark.parametrize( + 'hook', + [ + 'after_agent_run', + 'after_agent_call', + 'after_execution', + 'on_execution_event', + 'on_task_event', + 'on_push_sent', + ], +) +@pytest.mark.asyncio +async def test_after_and_observer_errors_are_swallowed(hook, caplog): + log: list[str] = [] + m = plugins.sequential(Recorder('a', log, fail={hook}), Recorder('b', log)) + calls = { + 'after_agent_run': lambda: m.after_agent_run( + ServerResponse('GetTask', None, ServerCallContext()) + ), + 'after_agent_call': lambda: m.after_agent_call( + ClientResponse('GetTask', None, ClientCallContext()) + ), + 'after_execution': lambda: m.after_execution(MagicMock(), None), + 'on_execution_event': lambda: m.on_execution_event(MagicMock(), Task()), + 'on_task_event': lambda: m.on_task_event(MagicMock()), + 'on_push_sent': lambda: m.on_push_sent(MagicMock()), + } + + await calls[hook]() + + assert f'a.{hook}' in log + assert f'b.{hook}' in log + assert f"A2A plugin 'a' failed in {hook}" in caplog.text + + +@pytest.mark.asyncio +async def test_before_push_send_error_propagates(): + m = plugins.sequential(Recorder('a', [], fail={'before_push_send'})) + with pytest.raises(ValueError, match='before_push_send'): + await m.before_push_send(push_request()) + + +def test_plugins_filtered_by_interface(): + class OnlyTask(TaskObserverPlugin): + pass + + class OnlyClient(ClientPlugin): + pass + + m = plugins.sequential(OnlyTask(), OnlyClient()) + + assert m.has_task_hooks + assert m.has_client_hooks + assert not m.has_server_hooks + assert not m.has_execution_hooks + assert not m.has_push_hooks + + +def test_empty_manager_has_no_hooks(): + m = SequentialPluginManager() + assert not any( + ( + m.has_server_hooks, + m.has_client_hooks, + m.has_execution_hooks, + m.has_task_hooks, + m.has_push_hooks, + ) + ) + assert m.plugins == () + + +def test_base_plugin_manager_has_no_hooks(): + m = plugins.PluginManager() + assert not m.has_server_hooks + assert not m.has_client_hooks + assert not m.has_execution_hooks + assert not m.has_task_hooks + assert not m.has_push_hooks + + +@pytest.mark.asyncio +async def test_nested_managers(): + log: list[str] = [] + inner = plugins.sequential(Recorder('b', log), Recorder('c', log)) + outer = plugins.sequential(Recorder('a', log), inner) + + await outer.before_agent_run(server_request()) + await outer.after_agent_run( + ServerResponse('GetTask', None, ServerCallContext()) + ) + + assert log == [ + 'a.before_agent_run', + 'b.before_agent_run', + 'c.before_agent_run', + 'c.after_agent_run', + 'b.after_agent_run', + 'a.after_agent_run', + ] + + +def test_nested_empty_manager_does_not_count_as_hooks(): + outer = plugins.sequential(plugins.sequential()) + assert not outer.has_server_hooks + assert not outer.has_push_hooks + + +def test_rejects_non_plugins(): + with pytest.raises(TypeError, match='Not A2A plugins'): + plugins.sequential(object()) # type: ignore[arg-type] + + +@pytest.mark.asyncio +async def test_base_hooks_are_no_ops(): + p = CompletePlugin() + assert await p.before_agent_run(server_request()) is None + assert ( + await p.after_agent_call( + ClientResponse('GetTask', None, ClientCallContext()) + ) + is None + ) + assert Plugin.name == 'unnamed' + e = ExecutionPlugin() + assert await e.on_execution_event(MagicMock(), Task()) is None + assert await PushPlugin().before_push_send(push_request()) is None + + +def test_global_registry(): + assert isinstance(plugins.get_global(), plugins.PluginManager) + m = plugins.sequential(Recorder('a', [])) + plugins.set_global(m) + assert plugins.get_global() is m + + +def test_set_global_requires_manager(): + with pytest.raises(TypeError, match='sequential'): + plugins.set_global(Recorder('a', [])) # type: ignore[arg-type] From 5b896c199b4d92dfbb8877e36663be4e18385eef Mon Sep 17 00:00:00 2001 From: Mykyta Netipa Date: Sat, 17 Oct 2026 15:17:46 +0000 Subject: [PATCH 02/11] feat(server): run plugin hooks in DefaultRequestHandlerV2 DefaultRequestHandlerV2 takes a plugins manager (defaulting to the global one) and runs: - before/after_agent_run around every public on_* method, with a call per stream event and a final call on completion, error or disconnect; - before/after_execution and on_execution_event around the agent's execute(), via a wrapping executor and event queue, so plugins see the detached execution and the time of each agent event; - on_task_event after every successful task write. TaskManager tracks the persisted state itself because it mutates the cached task in place before saving, which made the store's prev argument useless for from_state; - before_push_send/on_push_sent per webhook delivery, scoped to the V2 path through a contextvar so the shared push sender stays hook-free for LegacyRequestHandler. Transports record the protocol in ServerCallContext.state['a2a.transport'] and gRPC exposes invocation metadata as state['headers'], so plugins can read trace and tenant headers on every transport. With no plugins registered nothing is wrapped and behavior is unchanged. --- src/a2a/compat/v0_3/grpc_handler.py | 3 + src/a2a/compat/v0_3/jsonrpc_adapter.py | 6 +- src/a2a/compat/v0_3/rest_adapter.py | 8 + src/a2a/server/agent_execution/active_task.py | 13 +- .../agent_execution/active_task_registry.py | 5 + .../agent_execution/plugin_agent_executor.py | 83 +++ src/a2a/server/context.py | 3 + .../server/request_handlers/_plugin_hooks.py | 125 ++++ .../default_request_handler_v2.py | 55 +- .../server/request_handlers/grpc_handler.py | 18 +- src/a2a/server/routes/jsonrpc_dispatcher.py | 6 +- src/a2a/server/routes/rest_dispatcher.py | 6 +- .../tasks/base_push_notification_sender.py | 69 ++ src/a2a/server/tasks/task_manager.py | 32 +- .../test_plugin_agent_executor.py | 99 +++ ...test_default_request_handler_v2_plugins.py | 662 ++++++++++++++++++ .../request_handlers/test_grpc_handler.py | 25 + .../server/tasks/test_task_manager_plugins.py | 163 +++++ 18 files changed, 1372 insertions(+), 9 deletions(-) create mode 100644 src/a2a/server/agent_execution/plugin_agent_executor.py create mode 100644 src/a2a/server/request_handlers/_plugin_hooks.py create mode 100644 tests/server/agent_execution/test_plugin_agent_executor.py create mode 100644 tests/server/request_handlers/test_default_request_handler_v2_plugins.py create mode 100644 tests/server/tasks/test_task_manager_plugins.py diff --git a/src/a2a/compat/v0_3/grpc_handler.py b/src/a2a/compat/v0_3/grpc_handler.py index b7bec26ea..5f046c765 100644 --- a/src/a2a/compat/v0_3/grpc_handler.py +++ b/src/a2a/compat/v0_3/grpc_handler.py @@ -24,6 +24,7 @@ _ERROR_CODE_MAP, DefaultGrpcServerCallContextBuilder, GrpcServerCallContextBuilder, + set_grpc_call_metadata, ) from a2a.server.request_handlers.request_handler import RequestHandler from a2a.utils.errors import A2AError, InvalidParamsError @@ -64,6 +65,7 @@ async def _handle_unary( """Centralized error handling and context management for unary calls.""" try: server_context = self._context_builder.build(context) + set_grpc_call_metadata(server_context, context) result = await handler_func(server_context) except A2AError as e: await self.abort_context(e, context) @@ -79,6 +81,7 @@ async def _handle_stream( """Centralized error handling and context management for streaming calls.""" try: server_context = self._context_builder.build(context) + set_grpc_call_metadata(server_context, context) async for item in handler_func(server_context): yield item except A2AError as e: diff --git a/src/a2a/compat/v0_3/jsonrpc_adapter.py b/src/a2a/compat/v0_3/jsonrpc_adapter.py index 303f5f9f8..a0ed23762 100644 --- a/src/a2a/compat/v0_3/jsonrpc_adapter.py +++ b/src/a2a/compat/v0_3/jsonrpc_adapter.py @@ -26,7 +26,7 @@ from a2a.compat.v0_3 import types as types_v03 from a2a.compat.v0_3.context_builders import V03ServerCallContextBuilder from a2a.compat.v0_3.request_handler import RequestHandler03 -from a2a.server.context import ServerCallContext +from a2a.server.context import TRANSPORT_STATE_KEY, ServerCallContext from a2a.server.jsonrpc_models import ( InternalError as CoreInternalError, ) @@ -42,6 +42,7 @@ create_event_source_response, ) from a2a.utils import constants +from a2a.utils.constants import TransportProtocol from a2a.utils.version_validator import validate_version @@ -125,6 +126,9 @@ async def handle_request( ) call_context = self._context_builder.build(request) + call_context.state.setdefault( + TRANSPORT_STATE_KEY, TransportProtocol.JSONRPC.value + ) call_context.tenant = ( getattr(specific_request.params, 'tenant', '') if hasattr(specific_request, 'params') diff --git a/src/a2a/compat/v0_3/rest_adapter.py b/src/a2a/compat/v0_3/rest_adapter.py index 0c2449dfa..0bd618c44 100644 --- a/src/a2a/compat/v0_3/rest_adapter.py +++ b/src/a2a/compat/v0_3/rest_adapter.py @@ -32,12 +32,14 @@ from a2a.compat.v0_3.context_builders import V03ServerCallContextBuilder from a2a.compat.v0_3.rest_handler import REST03Handler +from a2a.server.context import TRANSPORT_STATE_KEY from a2a.server.routes.common import ( DefaultServerCallContextBuilder, ServerCallContextBuilder, create_event_source_response, ) from a2a.utils import json_utils +from a2a.utils.constants import TransportProtocol from a2a.utils.error_handlers import ( rest_error_handler, rest_stream_error_handler, @@ -75,6 +77,9 @@ async def _handle_request( request: Request, ) -> Response: call_context = self._context_builder.build(request) + call_context.state.setdefault( + TRANSPORT_STATE_KEY, TransportProtocol.HTTP_JSON.value + ) response = await method(request, call_context) return JSONResponse(content=response) @@ -92,6 +97,9 @@ async def _handle_streaming_request( ) from e call_context = self._context_builder.build(request) + call_context.state.setdefault( + TRANSPORT_STATE_KEY, TransportProtocol.HTTP_JSON.value + ) async def event_generator( stream: AsyncIterable[Any], diff --git a/src/a2a/server/agent_execution/active_task.py b/src/a2a/server/agent_execution/active_task.py index 5e3dab50f..bb6530b91 100644 --- a/src/a2a/server/agent_execution/active_task.py +++ b/src/a2a/server/agent_execution/active_task.py @@ -48,6 +48,7 @@ if TYPE_CHECKING: from collections.abc import AsyncGenerator, Callable + from a2a.plugins.manager import PluginManager from a2a.server.agent_execution.agent_executor import AgentExecutor from a2a.server.cluster.event_stream import TaskEventStream from a2a.server.context import ServerCallContext @@ -56,6 +57,7 @@ ) from a2a.server.tasks.task_manager import TaskManager +from a2a.plugins._runtime import push_plugins_scope from a2a.server.cluster.event_stream import VersionedEvent from a2a.server.cluster.task_store import ConcurrentTaskModificationError from a2a.server.events.event_queue_v2 import ( @@ -359,9 +361,10 @@ async def _update_task_state( self.active_task._task_id, ) try: - await self.active_task._push_sender.send_notification( - self.active_task._task_id, event - ) + with push_plugins_scope(self.active_task._plugins): + await self.active_task._push_sender.send_notification( + self.active_task._task_id, event + ) except Exception: # Push delivery is best-effort: a sender failure must not # fail the task or disturb the event stream. This guard is @@ -441,6 +444,7 @@ def __init__( # noqa: PLR0913 push_sender: PushNotificationSender | None = None, on_cleanup: Callable[[ActiveTask], None] | None = None, event_stream: TaskEventStream | None = None, + plugins: PluginManager | None = None, ) -> None: """Initializes the ActiveTask. @@ -454,6 +458,8 @@ def __init__( # noqa: PLR0913 the task from the ActiveTaskRegistry. event_stream: Optional cross-replica stream; applied events are published to it so other replicas can observe them. + plugins: Optional plugin manager whose push hooks run when + `push_sender` delivers notifications. """ # --- Core Dependencies --- self._agent_executor = agent_executor @@ -466,6 +472,7 @@ def __init__( # noqa: PLR0913 self._push_sender = push_sender self._on_cleanup = on_cleanup self._event_stream = event_stream + self._plugins = plugins # --- Synchronization Primitives --- # `_lock` protects structural lifecycle changes: start(), subscribe() counting, diff --git a/src/a2a/server/agent_execution/active_task_registry.py b/src/a2a/server/agent_execution/active_task_registry.py index cae2a55ce..11441aa35 100644 --- a/src/a2a/server/agent_execution/active_task_registry.py +++ b/src/a2a/server/agent_execution/active_task_registry.py @@ -8,6 +8,7 @@ if TYPE_CHECKING: + from a2a.plugins.manager import PluginManager from a2a.server.agent_execution.agent_executor import AgentExecutor from a2a.server.cluster.event_stream import TaskEventStream from a2a.server.cluster.task_store import VersionedTaskStore @@ -34,8 +35,10 @@ def __init__( task_store: TaskStore | VersionedTaskStore, push_sender: PushNotificationSender | None = None, event_stream: TaskEventStream | None = None, + plugins: PluginManager | None = None, ): self._agent_executor = agent_executor + self._plugins = plugins self._task_store = task_store self._push_sender = push_sender self._event_stream = event_stream @@ -64,6 +67,7 @@ async def get_or_create( task_store=self._task_store, initial_message=initial_message, context=call_context, + plugins=self._plugins, ) active_task = ActiveTask( @@ -73,6 +77,7 @@ async def get_or_create( push_sender=self._push_sender, on_cleanup=self._on_active_task_cleanup, event_stream=self._event_stream, + plugins=self._plugins, ) self._active_tasks[task_id] = active_task diff --git a/src/a2a/server/agent_execution/plugin_agent_executor.py b/src/a2a/server/agent_execution/plugin_agent_executor.py new file mode 100644 index 000000000..198611230 --- /dev/null +++ b/src/a2a/server/agent_execution/plugin_agent_executor.py @@ -0,0 +1,83 @@ +"""Runs `ExecutionPlugin` hooks around an `AgentExecutor`.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +from a2a.server.agent_execution.agent_executor import AgentExecutor +from a2a.server.events.event_queue import EventQueue + + +if TYPE_CHECKING: + from a2a.plugins.manager import PluginManager + from a2a.server.agent_execution.context import RequestContext + from a2a.server.events.event_queue import Event + + +class PluginEventQueue(EventQueue): + """Calls `on_execution_event` for each event the agent enqueues. + + Wraps the queue passed to `AgentExecutor.execute`. Events the SDK enqueues + on the underlying queue itself bypass the wrapper and are not reported. + """ + + def __init__( + self, + inner: EventQueue, + plugins: PluginManager, + request_context: RequestContext, + ) -> None: + self._inner = inner + self._plugins = plugins + self._request_context = request_context + + async def enqueue_event(self, event: Event) -> None: + """Reports the event to plugins, then enqueues it.""" + await self._plugins.on_execution_event(self._request_context, event) + await self._inner.enqueue_event(event) + + def __getattr__(self, name: str) -> Any: + """Forwards other attributes to the wrapped queue. + + Keeps executors working that rely on the concrete queue type. + """ + if name == '_inner': + raise AttributeError(name) + return getattr(self._inner, name) + + +class PluginAgentExecutor(AgentExecutor): + """Runs `before_execution` / `after_execution` around ``execute``.""" + + def __init__(self, inner: AgentExecutor, plugins: PluginManager) -> None: + self._inner = inner + self._plugins = plugins + + @property + def inner(self) -> AgentExecutor: + """The wrapped executor.""" + return self._inner + + async def execute( + self, context: RequestContext, event_queue: EventQueue + ) -> None: + """Runs the wrapped executor with execution hooks.""" + # If a before-hook raises, the manager unwinds the plugins that ran + # and the exception fails the execution. + await self._plugins.before_execution(context) + error: BaseException | None = None + try: + await self._inner.execute( + context, PluginEventQueue(event_queue, self._plugins, context) + ) + except Exception as e: + error = e + raise + finally: + await self._plugins.after_execution(context, error) + + async def cancel( + self, context: RequestContext, event_queue: EventQueue + ) -> None: + """Delegates to the wrapped executor without hooks.""" + await self._inner.cancel(context, event_queue) diff --git a/src/a2a/server/context.py b/src/a2a/server/context.py index 833ca44c4..c9059014f 100644 --- a/src/a2a/server/context.py +++ b/src/a2a/server/context.py @@ -10,6 +10,9 @@ State = collections.abc.MutableMapping[str, typing.Any] +TRANSPORT_STATE_KEY = 'a2a.transport' +"""`ServerCallContext.state` key holding the transport of the request.""" + class ServerCallContext(BaseModel): """A context passed when calling a server method. diff --git a/src/a2a/server/request_handlers/_plugin_hooks.py b/src/a2a/server/request_handlers/_plugin_hooks.py new file mode 100644 index 000000000..5774e08b9 --- /dev/null +++ b/src/a2a/server/request_handlers/_plugin_hooks.py @@ -0,0 +1,125 @@ +"""Runs `ServerPlugin` hooks around `DefaultRequestHandlerV2` methods.""" + +from __future__ import annotations + +import functools +import inspect + +from contextlib import aclosing +from typing import TYPE_CHECKING, Any, Protocol, TypeVar, cast + +from a2a.plugins.types import ServerRequest, ServerResponse + + +if TYPE_CHECKING: + from collections.abc import AsyncGenerator, Callable + + from a2a.plugins.manager import PluginManager + from a2a.server.context import ServerCallContext + from a2a.types.a2a_pb2 import AgentCard + + +class _PluginHost(Protocol): + """A request handler that exposes its plugin manager and agent card.""" + + @property + def plugins(self) -> PluginManager: ... + + @property + def agent_card(self) -> AgentCard: ... + + +_F = TypeVar('_F', bound='Callable[..., Any]') + + +def server_plugin_hooks(method: str) -> Callable[[_F], _F]: + """Runs `before_agent_run` / `after_agent_run` around a handler method. + + Apply as the outermost decorator, so that validation errors raised by + inner decorators also reach `after_agent_run`. When the handler's manager + has no server hooks the method runs unchanged. + + Args: + method: The A2A RPC name passed to the hooks, e.g. ``'GetTask'``. + """ + + def decorator(func: _F) -> _F: + if inspect.isasyncgenfunction(func): + + @functools.wraps(func) + async def stream_wrapper( + self: _PluginHost, + params: Any, + context: ServerCallContext, + *args: Any, + **kwargs: Any, + ) -> AsyncGenerator[Any, None]: + plugins = self.plugins + if not plugins.has_server_hooks: + async with aclosing( + func(self, params, context, *args, **kwargs) + ) as events: + async for event in events: + yield event + return + + await plugins.before_agent_run( + ServerRequest(method, params, context, self.agent_card) + ) + error: BaseException | None = None + try: + async with aclosing( + func(self, params, context, *args, **kwargs) + ) as events: + async for event in events: + await plugins.after_agent_run( + ServerResponse( + method, + event, + context, + is_stream_event=True, + is_final=False, + ) + ) + yield event + except Exception as e: + error = e + raise + finally: + # Also runs when the client disconnects (aclose()). + await plugins.after_agent_run( + ServerResponse(method, None, context, error=error) + ) + + return cast('_F', stream_wrapper) + + @functools.wraps(func) + async def unary_wrapper( + self: _PluginHost, + params: Any, + context: ServerCallContext, + *args: Any, + **kwargs: Any, + ) -> Any: + plugins = self.plugins + if not plugins.has_server_hooks: + return await func(self, params, context, *args, **kwargs) + + await plugins.before_agent_run( + ServerRequest(method, params, context, self.agent_card) + ) + try: + result = await func(self, params, context, *args, **kwargs) + except Exception as e: + await plugins.after_agent_run( + ServerResponse(method, None, context, error=e) + ) + raise + await plugins.after_agent_run( + ServerResponse(method, result, context) + ) + return result + + return cast('_F', unary_wrapper) + + return decorator diff --git a/src/a2a/server/request_handlers/default_request_handler_v2.py b/src/a2a/server/request_handlers/default_request_handler_v2.py index b135f529f..16422da49 100644 --- a/src/a2a/server/request_handlers/default_request_handler_v2.py +++ b/src/a2a/server/request_handlers/default_request_handler_v2.py @@ -7,6 +7,8 @@ from contextlib import aclosing from typing import TYPE_CHECKING, Any, cast +from a2a.plugins.manager import get_global +from a2a.plugins.types import TaskEventInfo from a2a.server.agent_execution import ( AgentExecutor, RequestContext, @@ -18,12 +20,16 @@ TERMINAL_TASK_STATES, ) from a2a.server.agent_execution.active_task_registry import ActiveTaskRegistry +from a2a.server.agent_execution.plugin_agent_executor import ( + PluginAgentExecutor, +) from a2a.server.cluster.event_stream import VersionedEvent from a2a.server.cluster.task_store import ( ConcurrentTaskModificationError, LegacyTaskStoreAdapter, VersionedTaskStore, ) +from a2a.server.request_handlers._plugin_hooks import server_plugin_hooks from a2a.server.request_handlers.request_handler import ( RequestHandler, validate, @@ -73,6 +79,7 @@ if TYPE_CHECKING: from collections.abc import AsyncGenerator, Awaitable, Callable + from a2a.plugins.manager import PluginManager from a2a.server.agent_execution.active_task import ActiveTask from a2a.server.cluster.event_stream import TaskEventStream from a2a.server.cluster.version import TaskVersion @@ -96,6 +103,10 @@ class DefaultRequestHandlerV2(RequestHandler): with `DefaultRequestHandler` but is not used: v2 delegates event streaming to an in-memory `ActiveTaskRegistry`. Passing a non-``None`` value emits a `DeprecationWarning` and logs a warning. + + ``plugins`` selects the plugin manager whose hooks run for this handler. + It defaults to the global manager (`a2a.plugins.get_global()`) at + construction time. """ _background_tasks: set[asyncio.Task] @@ -118,6 +129,7 @@ def __init__( # noqa: PLR0913 push_url_validator: Callable[[str], Awaitable[bool]] | None = None, validate_input_modes: bool = False, event_stream: TaskEventStream | None = None, + plugins: PluginManager | None = None, ) -> None: if queue_manager is not None: message = ( @@ -130,6 +142,9 @@ def __init__( # noqa: PLR0913 ) warnings.warn(message, DeprecationWarning, stacklevel=2) logger.warning(message) + self._plugins: PluginManager = ( + plugins if plugins is not None else get_global() + ) self.agent_executor = agent_executor self.task_store: TaskStore | VersionedTaskStore = task_store self._versioned_store: VersionedTaskStore = ( @@ -168,13 +183,28 @@ def __init__( # noqa: PLR0913 ) ) self._active_task_registry = ActiveTaskRegistry( - agent_executor=self.agent_executor, + agent_executor=( + PluginAgentExecutor(agent_executor, self._plugins) + if self._plugins.has_execution_hooks + else agent_executor + ), task_store=self._versioned_store, push_sender=self._push_sender, event_stream=self._event_stream, + plugins=self._plugins, ) self._background_tasks = set() + @property + def plugins(self) -> PluginManager: + """The plugin manager whose hooks run for this handler.""" + return self._plugins + + @property + def agent_card(self) -> AgentCard: + """The agent card served by this handler.""" + return self._agent_card + async def _reject_unsafe_push_url(self, url: str) -> None: """Apply the configured push-URL policy, if any.""" if self._push_url_validator is None: @@ -191,6 +221,7 @@ async def aclose(self) -> None: """ await self._active_task_registry.aclose() + @server_plugin_hooks('GetTask') @validate_request_params async def on_get_task( # noqa: D102 self, @@ -206,6 +237,7 @@ async def on_get_task( # noqa: D102 return apply_history_length(stored.task, params) + @server_plugin_hooks('ListTasks') @validate_request_params async def on_list_tasks( # noqa: D102 self, @@ -227,6 +259,7 @@ async def on_list_tasks( # noqa: D102 return page + @server_plugin_hooks('CancelTask') @validate_request_params async def on_cancel_task( # noqa: D102 self, @@ -313,6 +346,18 @@ async def _write_cancel( prev_version=version, context=context, ) + if self._plugins.has_task_hooks: + # Remote cancel writes the store directly, bypassing TaskManager. + await self._plugins.on_task_event( + TaskEventInfo( + task_id=task_id, + from_state=current.status.state, + to_state=cancelled.status.state, + event=event, + task=cancelled, + context=context, + ) + ) if self._event_stream is not None: await self._event_stream.publish( task_id, VersionedEvent(event=event, version=new_version) @@ -392,6 +437,7 @@ async def _setup_active_task( return active_task, request_context + @server_plugin_hooks('SendMessage') @validate_request_params async def on_message_send( # noqa: D102 self, @@ -457,6 +503,7 @@ async def on_message_send( # noqa: D102 ) return result + @server_plugin_hooks('SendStreamingMessage') @validate_request_params @validate( lambda self: self._agent_card.capabilities.streaming, @@ -489,6 +536,7 @@ async def on_message_send_stream( # noqa: D102 else: yield event + @server_plugin_hooks('CreateTaskPushNotificationConfig') @validate_request_params @validate( lambda self: self._agent_card.capabilities.push_notifications, @@ -519,6 +567,7 @@ async def on_create_task_push_notification_config( # noqa: D102 # Custom stores written before set_info returned the stored config. return normalize_push_notification_config(task_id, params) + @server_plugin_hooks('GetTaskPushNotificationConfig') @validate_request_params @validate( lambda self: self._agent_card.capabilities.push_notifications, @@ -548,6 +597,7 @@ async def on_get_task_push_notification_config( # noqa: D102 raise TaskNotFoundError + @server_plugin_hooks('SubscribeToTask') @validate_request_params @validate( lambda self: self._agent_card.capabilities.streaming, @@ -622,6 +672,7 @@ async def _subscribe_remote( ): return + @server_plugin_hooks('ListTaskPushNotificationConfigs') @validate_request_params @validate( lambda self: self._agent_card.capabilities.push_notifications, @@ -648,6 +699,7 @@ async def on_list_task_push_notification_configs( # noqa: D102 configs=push_notification_config_list ) + @server_plugin_hooks('DeleteTaskPushNotificationConfig') @validate_request_params @validate( lambda self: self._agent_card.capabilities.push_notifications, @@ -669,6 +721,7 @@ async def on_delete_task_push_notification_config( # noqa: D102 await self._push_config_store.delete_info(task_id, context, config_id) + @server_plugin_hooks('GetExtendedAgentCard') @validate_request_params @validate( lambda self: self._agent_card.capabilities.extended_agent_card, diff --git a/src/a2a/server/request_handlers/grpc_handler.py b/src/a2a/server/request_handlers/grpc_handler.py index d4d8ed669..bfbcbb6c0 100644 --- a/src/a2a/server/request_handlers/grpc_handler.py +++ b/src/a2a/server/request_handlers/grpc_handler.py @@ -27,10 +27,11 @@ HTTP_EXTENSION_HEADER, get_requested_extensions, ) -from a2a.server.context import ServerCallContext +from a2a.server.context import TRANSPORT_STATE_KEY, ServerCallContext from a2a.server.request_handlers.request_handler import RequestHandler from a2a.types import a2a_pb2 from a2a.utils import proto_utils +from a2a.utils.constants import TransportProtocol from a2a.utils.errors import A2A_ERROR_REASONS, A2AError, TaskNotFoundError from a2a.utils.grpc_status import status_to_grpc from a2a.utils.proto_utils import validation_errors_to_bad_request @@ -81,6 +82,20 @@ def _get_metadata_value( ] +def set_grpc_call_metadata( + server_context: ServerCallContext, context: grpc.aio.ServicerContext +) -> None: + """Records the transport and text invocation metadata in the call state.""" + server_context.state.setdefault( + TRANSPORT_STATE_KEY, TransportProtocol.GRPC.value + ) + if 'headers' not in server_context.state: + md = context.invocation_metadata() or () + server_context.state['headers'] = { + k.lower(): v for k, v in md if isinstance(v, str) + } + + _ERROR_CODE_MAP: dict[type[A2AError], grpc.StatusCode] = { types.InvalidRequestError: grpc.StatusCode.INVALID_ARGUMENT, types.MethodNotFoundError: grpc.StatusCode.NOT_FOUND, @@ -424,5 +439,6 @@ def _build_call_context( request: message.Message, ) -> ServerCallContext: server_context = self._context_builder.build(context) + set_grpc_call_metadata(server_context, context) server_context.tenant = getattr(request, 'tenant', '') return server_context diff --git a/src/a2a/server/routes/jsonrpc_dispatcher.py b/src/a2a/server/routes/jsonrpc_dispatcher.py index 21b71e921..b4cbaa241 100644 --- a/src/a2a/server/routes/jsonrpc_dispatcher.py +++ b/src/a2a/server/routes/jsonrpc_dispatcher.py @@ -10,7 +10,7 @@ from jsonrpc.jsonrpc2 import JSONRPC20Request, JSONRPC20Response from a2a.compat.v0_3.jsonrpc_adapter import JSONRPC03Adapter -from a2a.server.context import ServerCallContext +from a2a.server.context import TRANSPORT_STATE_KEY, ServerCallContext from a2a.server.events import Event from a2a.server.jsonrpc_models import ( InternalError, @@ -41,6 +41,7 @@ TaskPushNotificationConfig, ) from a2a.utils import constants, json_utils, proto_utils +from a2a.utils.constants import TransportProtocol from a2a.utils.errors import ( A2AError, InvalidParamsError, @@ -308,6 +309,9 @@ async def handle_requests(self, request: Request) -> Response: # noqa: PLR0911, # 3) Build call context and wrap the request for downstream handling call_context = self._context_builder.build(request) + call_context.state.setdefault( + TRANSPORT_STATE_KEY, TransportProtocol.JSONRPC.value + ) call_context.tenant = getattr(specific_request, 'tenant', '') call_context.state['method'] = method call_context.state['request_id'] = request_id diff --git a/src/a2a/server/routes/rest_dispatcher.py b/src/a2a/server/routes/rest_dispatcher.py index 54150396e..3000b8cf7 100644 --- a/src/a2a/server/routes/rest_dispatcher.py +++ b/src/a2a/server/routes/rest_dispatcher.py @@ -5,7 +5,7 @@ from google.protobuf.json_format import MessageToDict, Parse -from a2a.server.context import ServerCallContext +from a2a.server.context import TRANSPORT_STATE_KEY, ServerCallContext from a2a.server.request_handlers.request_handler import RequestHandler from a2a.server.routes.common import ( DefaultServerCallContextBuilder, @@ -20,6 +20,7 @@ SubscribeToTaskRequest, ) from a2a.utils import constants, json_utils, proto_utils +from a2a.utils.constants import TransportProtocol from a2a.utils.error_handlers import ( build_rest_error_payload, rest_error_handler, @@ -106,6 +107,9 @@ def __init__( def _build_call_context(self, request: Request) -> ServerCallContext: call_context = self._context_builder.build(request) + call_context.state.setdefault( + TRANSPORT_STATE_KEY, TransportProtocol.HTTP_JSON.value + ) if 'tenant' in request.path_params: call_context.tenant = request.path_params['tenant'] return call_context diff --git a/src/a2a/server/tasks/base_push_notification_sender.py b/src/a2a/server/tasks/base_push_notification_sender.py index 070361a52..26c9c0852 100644 --- a/src/a2a/server/tasks/base_push_notification_sender.py +++ b/src/a2a/server/tasks/base_push_notification_sender.py @@ -1,5 +1,6 @@ import asyncio import logging +import time from collections.abc import Awaitable, Callable @@ -7,6 +8,9 @@ from google.protobuf.json_format import MessageToDict +from a2a.plugins._runtime import current_push_plugins +from a2a.plugins.manager import PluginManager +from a2a.plugins.types import PushRequest, PushResult from a2a.server.context import ServerCallContext from a2a.server.tasks.push_notification_config_store import ( PushNotificationConfigStore, @@ -121,6 +125,7 @@ async def _dispatch_notification( return False if not accepted: return False + plugins = current_push_plugins() try: headers: dict[str, str] = {} if push_info.token: @@ -140,6 +145,10 @@ async def _dispatch_notification( task_id, ) + if plugins is not None: + return await self._post_with_plugins( + plugins, PushRequest(task_id, url, headers, event) + ) response = await self._client.post( url, json=MessageToDict(to_stream_response(event)), @@ -157,3 +166,63 @@ async def _dispatch_notification( ) return False return True + + async def _post_with_plugins( + self, plugins: PluginManager, request: PushRequest + ) -> bool: + """Delivers one notification with `PushPlugin` hooks. + + A raising `before_push_send` fails the delivery. The webhook URL + cannot be changed by plugins, since it was validated already. + """ + url = request.url + try: + await plugins.before_push_send(request) + except Exception as e: + logger.exception( + 'Push plugin failed before sending for task_id=%s to URL: %s.', + request.task_id, + request.url, + ) + await plugins.on_push_sent(PushResult(request, 0.0, None, e)) + return False + if request.skip: + logger.debug( + 'Push-notification for task_id=%s to URL: %s skipped by a ' + 'plugin', + request.task_id, + request.url, + ) + await plugins.on_push_sent( + PushResult(request, 0.0, None, None, skipped=True) + ) + return True + + start = time.monotonic() + status_code: int | None = None + error: Exception | None = None + try: + response = await self._client.post( + url, + json=MessageToDict(to_stream_response(request.event)), + headers=request.headers, + ) + status_code = response.status_code + response.raise_for_status() + except Exception as e: + error = e + logger.exception( + 'Error sending push-notification for task_id=%s to URL: %s.', + request.task_id, + request.url, + ) + await plugins.on_push_sent( + PushResult(request, time.monotonic() - start, status_code, error) + ) + if error is None: + logger.info( + 'Push-notification sent for task_id=%s to URL: %s', + request.task_id, + request.url, + ) + return error is None diff --git a/src/a2a/server/tasks/task_manager.py b/src/a2a/server/tasks/task_manager.py index 152083459..420f9f239 100644 --- a/src/a2a/server/tasks/task_manager.py +++ b/src/a2a/server/tasks/task_manager.py @@ -4,6 +4,7 @@ from typing import TYPE_CHECKING +from a2a.plugins.types import TaskEventInfo from a2a.types.a2a_pb2 import ( Artifact, Message, @@ -18,6 +19,7 @@ if TYPE_CHECKING: + from a2a.plugins.manager import PluginManager from a2a.server.cluster.task_store import VersionedTaskStore from a2a.server.cluster.version import TaskVersion from a2a.server.context import ServerCallContext @@ -103,13 +105,14 @@ class TaskManager: events received from the agent. """ - def __init__( + def __init__( # noqa: PLR0913 self, task_store: TaskStore | VersionedTaskStore, context: ServerCallContext, task_id: str | None, context_id: str | None, initial_message: Message | None, + plugins: PluginManager | None = None, ): """Initializes the TaskManager. @@ -121,6 +124,8 @@ def __init__( context_id: The ID of the context, if known from the request. initial_message: The `Message` that initiated the task, if any. Used when creating a new task object. + plugins: Plugin manager notified of every successful task write + via `on_task_event`. ``None`` disables the notifications. """ if task_id is not None and not (isinstance(task_id, str) and task_id): raise ValueError('Task ID must be a non-empty string') @@ -149,6 +154,13 @@ def __init__( self._current_version: TaskVersion = TaskVersion.MISSING # Transient: the event currently being persisted, read by _save_task. self._pending_event: Event | None = None + self._plugins = ( + plugins if plugins is not None and plugins.has_task_hooks else None + ) + # State of the task as last read from or written to the store, None + # if it is not known to exist. Tracked separately because writes + # mutate `_current_task` in place before saving it. + self._persisted_state: TaskState | None = None logger.debug( 'TaskManager initialized with task_id: %s, context_id: %s', task_id, @@ -165,6 +177,7 @@ def invalidate(self) -> None: self._current_task = None self._current_version = TaskVersion.MISSING + self._persisted_state = None @property def current_version(self) -> TaskVersion: @@ -198,10 +211,12 @@ async def get_task(self) -> Task | None: if stored is not None: self._current_task = stored.task self._current_version = stored.version + self._persisted_state = stored.task.status.state logger.debug('Task %s retrieved successfully.', self.task_id) else: self._current_task = None self._current_version = TaskVersion.MISSING + self._persisted_state = None logger.debug('Task %s not found.', self.task_id) return self._current_task @@ -297,8 +312,10 @@ async def ensure_task_id(self, task_id: str, context_id: str) -> Task: if stored is not None: task = stored.task self._current_version = stored.version + self._persisted_state = stored.task.status.state else: self._current_version = TaskVersion.MISSING + self._persisted_state = None if not task: logger.info( @@ -388,6 +405,19 @@ async def _save_task(self, task: Task) -> None: logger.info('New task created with id: %s', task.id) self.task_id = task.id self.context_id = task.context_id + from_state = self._persisted_state + self._persisted_state = task.status.state + if self._plugins is not None: + await self._plugins.on_task_event( + TaskEventInfo( + task_id=task.id, + from_state=from_state, + to_state=task.status.state, + event=self._pending_event, + task=task, + context=self._call_context, + ) + ) def update_with_message(self, message: Message, task: Task) -> Task: """Updates a task object in memory by adding a new message to its history. diff --git a/tests/server/agent_execution/test_plugin_agent_executor.py b/tests/server/agent_execution/test_plugin_agent_executor.py new file mode 100644 index 000000000..91e8ed631 --- /dev/null +++ b/tests/server/agent_execution/test_plugin_agent_executor.py @@ -0,0 +1,99 @@ +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from a2a import plugins +from a2a.server.agent_execution import AgentExecutor +from a2a.server.agent_execution.plugin_agent_executor import ( + PluginAgentExecutor, + PluginEventQueue, +) +from a2a.server.events import EventQueue +from a2a.types.a2a_pb2 import Message, Part, Role + + +class InnerQueue(EventQueue): + def __init__(self) -> None: + self.events: list = [] + self.extra = 'forwarded' + + async def enqueue_event(self, event) -> None: + self.events.append(event) + + +class Collect(plugins.ExecutionPlugin): + def __init__(self) -> None: + self.log: list = [] + + async def before_execution(self, request_context) -> None: + self.log.append('before') + + async def on_execution_event(self, request_context, event) -> None: + self.log.append(('event', event.message_id)) + + async def after_execution(self, request_context, error) -> None: + self.log.append(('after', error)) + + +def message(message_id: str) -> Message: + return Message( + role=Role.ROLE_AGENT, message_id=message_id, parts=[Part(text='x')] + ) + + +@pytest.mark.asyncio +async def test_event_queue_reports_then_enqueues_and_forwards_attributes(): + collect = Collect() + inner = InnerQueue() + queue = PluginEventQueue(inner, plugins.sequential(collect), MagicMock()) + + await queue.enqueue_event(message('m-1')) + + assert collect.log == [('event', 'm-1')] + assert [e.message_id for e in inner.events] == ['m-1'] + assert queue.extra == 'forwarded' + with pytest.raises(AttributeError): + _ = queue.missing_attribute + + +@pytest.mark.asyncio +async def test_executor_wraps_execute_and_passes_errors(): + collect = Collect() + + class Agent(AgentExecutor): + async def execute(self, context, event_queue): + await event_queue.enqueue_event(message('m-1')) + if context == 'fail': + raise ValueError('boom') + + async def cancel(self, context, event_queue): + pass + + agent = Agent() + executor = PluginAgentExecutor(agent, plugins.sequential(collect)) + assert executor.inner is agent + + inner = InnerQueue() + await executor.execute(MagicMock(), inner) + assert collect.log == ['before', ('event', 'm-1'), ('after', None)] + assert len(inner.events) == 1 + + collect.log.clear() + with pytest.raises(ValueError, match='boom'): + await executor.execute('fail', InnerQueue()) + assert collect.log[-1][0] == 'after' + assert isinstance(collect.log[-1][1], ValueError) + + +@pytest.mark.asyncio +async def test_cancel_delegates_without_hooks(): + collect = Collect() + agent = MagicMock(spec=AgentExecutor) + agent.cancel = AsyncMock() + executor = PluginAgentExecutor(agent, plugins.sequential(collect)) + queue = InnerQueue() + + await executor.cancel('ctx', queue) + + agent.cancel.assert_awaited_once_with('ctx', queue) + assert collect.log == [] diff --git a/tests/server/request_handlers/test_default_request_handler_v2_plugins.py b/tests/server/request_handlers/test_default_request_handler_v2_plugins.py new file mode 100644 index 000000000..1e59d3972 --- /dev/null +++ b/tests/server/request_handlers/test_default_request_handler_v2_plugins.py @@ -0,0 +1,662 @@ +"""Plugin hooks run by DefaultRequestHandlerV2 (request, execution, task, push).""" + +import asyncio + +from collections.abc import Callable +from typing import Any + +import httpx +import pytest + +from a2a import plugins +from a2a.helpers.proto_helpers import new_task_from_user_message +from a2a.plugins import ( + CompletePlugin, + ExecutionPlugin, + PushPlugin, + PushRequest, + PushResult, + ServerRequest, + ServerResponse, + TaskEventInfo, + TaskObserverPlugin, +) +from a2a.server.agent_execution import AgentExecutor, RequestContext +from a2a.server.agent_execution.plugin_agent_executor import ( + PluginAgentExecutor, +) +from a2a.server.context import ServerCallContext +from a2a.server.events import EventQueue +from a2a.server.request_handlers import ( + DefaultRequestHandlerV2, + LegacyRequestHandler, +) +from a2a.server.tasks import ( + InMemoryPushNotificationConfigStore, + InMemoryTaskStore, + TaskUpdater, +) +from a2a.server.tasks.base_push_notification_sender import ( + BasePushNotificationSender, +) +from a2a.types import TaskNotFoundError +from a2a.types.a2a_pb2 import ( + AgentCapabilities, + AgentCard, + CancelTaskRequest, + GetTaskRequest, + Message, + Part, + Role, + SendMessageConfiguration, + SendMessageRequest, + Task, + TaskPushNotificationConfig, + TaskState, + TaskStatus, +) + + +CARD = AgentCard( + name='plugin_test_agent', + version='1.0', + capabilities=AgentCapabilities(streaming=True, push_notifications=True), +) + + +class Recorder(CompletePlugin, ExecutionPlugin, TaskObserverPlugin, PushPlugin): + """Records (hook, argument) for every hook call.""" + + name = 'recorder' + + def __init__(self) -> None: + self.calls: list[tuple[str, Any]] = [] + + def hooks(self, name: str) -> list[Any]: + return [arg for hook, arg in self.calls if hook == name] + + def names(self) -> list[str]: + return [hook for hook, _ in self.calls] + + async def before_agent_run(self, request: ServerRequest) -> None: + self.calls.append(('before_agent_run', request)) + + async def after_agent_run(self, response: ServerResponse) -> None: + self.calls.append(('after_agent_run', response)) + + async def before_execution(self, request_context: RequestContext) -> None: + self.calls.append(('before_execution', request_context)) + + async def on_execution_event(self, request_context, event) -> None: + self.calls.append(('on_execution_event', event)) + + async def after_execution(self, request_context, error) -> None: + self.calls.append(('after_execution', error)) + + async def on_task_event(self, info: TaskEventInfo) -> None: + self.calls.append(('on_task_event', info)) + + async def before_push_send(self, request: PushRequest) -> None: + request.headers['x-plugin'] = 'yes' + self.calls.append(('before_push_send', request)) + + async def on_push_sent(self, result: PushResult) -> None: + self.calls.append(('on_push_sent', result)) + + +class TaskAgent(AgentExecutor): + """SUBMITTED -> WORKING -> artifact -> COMPLETED.""" + + async def execute(self, context: RequestContext, event_queue: EventQueue): + task = context.current_task or new_task_from_user_message( + context.message + ) + if not context.current_task: + await event_queue.enqueue_event(task) + updater = TaskUpdater(event_queue, task.id, task.context_id) + await updater.start_work() + await updater.add_artifact(parts=[Part(text='result')], name='out') + await updater.complete() + + async def cancel(self, context: RequestContext, event_queue: EventQueue): + pass + + +class MessageAgent(AgentExecutor): + """Replies with a single Message, creating no task.""" + + async def execute(self, context: RequestContext, event_queue: EventQueue): + await event_queue.enqueue_event( + Message( + role=Role.ROLE_AGENT, + message_id='reply-1', + parts=[Part(text='hi')], + ) + ) + + async def cancel(self, context: RequestContext, event_queue: EventQueue): + pass + + +def send_request( + text: str = 'hello', + configuration: SendMessageConfiguration | None = None, +) -> SendMessageRequest: + request = SendMessageRequest( + message=Message( + role=Role.ROLE_USER, message_id='m-1', parts=[Part(text=text)] + ) + ) + if configuration is not None: + request.configuration.CopyFrom(configuration) + return request + + +async def wait_for(predicate: Callable[[], bool], timeout: float = 5) -> None: + async def poll() -> None: + while not predicate(): + await asyncio.sleep(0.01) + + await asyncio.wait_for(poll(), timeout) + + +def make_handler( + recorder: Recorder | None, + executor: AgentExecutor | None = None, + **kwargs: Any, +) -> DefaultRequestHandlerV2: + return DefaultRequestHandlerV2( + agent_executor=executor or TaskAgent(), + task_store=kwargs.pop('task_store', InMemoryTaskStore()), + agent_card=CARD, + plugins=plugins.sequential(recorder) if recorder else None, + **kwargs, + ) + + +@pytest.fixture(autouse=True) +def restore_global(): + previous = plugins.get_global() + yield + plugins.set_global(previous) + + +def test_no_plugins_leaves_executor_unwrapped(): + executor = TaskAgent() + handler = DefaultRequestHandlerV2( + agent_executor=executor, task_store=InMemoryTaskStore(), agent_card=CARD + ) + + assert not handler.plugins.has_server_hooks + assert handler._active_task_registry._agent_executor is executor + assert handler.agent_card is CARD + + +def test_global_manager_is_used_by_default(): + recorder = Recorder() + manager = plugins.sequential(recorder) + plugins.set_global(manager) + executor = TaskAgent() + + handler = DefaultRequestHandlerV2( + agent_executor=executor, task_store=InMemoryTaskStore(), agent_card=CARD + ) + + assert handler.plugins is manager + assert handler.agent_executor is executor + assert isinstance( + handler._active_task_registry._agent_executor, PluginAgentExecutor + ) + + +@pytest.mark.asyncio +async def test_no_plugins_request_still_works(): + handler = make_handler(None) + result = await handler.on_message_send(send_request(), ServerCallContext()) + assert result.status.state == TaskState.TASK_STATE_COMPLETED + + +@pytest.mark.asyncio +async def test_send_message_runs_all_server_hooks(): + recorder = Recorder() + handler = make_handler(recorder) + context = ServerCallContext() + + result = await handler.on_message_send(send_request(), context) + await wait_for(lambda: 'after_execution' in recorder.names()) + + assert isinstance(result, Task) + assert result.status.state == TaskState.TASK_STATE_COMPLETED + + (request,) = recorder.hooks('before_agent_run') + assert request.method == 'SendMessage' + assert request.context is context + assert request.agent_card is CARD + assert request.payload.message.message_id == 'm-1' + + (response,) = recorder.hooks('after_agent_run') + assert response.method == 'SendMessage' + assert response.payload is result + assert response.error is None + assert response.is_final + assert not response.is_stream_event + + # The execution runs with the same call context as the request. + (request_context,) = recorder.hooks('before_execution') + assert request_context.call_context is context + assert recorder.hooks('after_execution') == [None] + + events = recorder.hooks('on_execution_event') + assert [type(e).__name__ for e in events] == [ + 'Task', + 'TaskStatusUpdateEvent', + 'TaskArtifactUpdateEvent', + 'TaskStatusUpdateEvent', + ] + + infos = recorder.hooks('on_task_event') + transitions = [(i.from_state, i.to_state) for i in infos] + assert transitions[0] == (None, TaskState.TASK_STATE_SUBMITTED) + assert ( + TaskState.TASK_STATE_SUBMITTED, + TaskState.TASK_STATE_WORKING, + ) in transitions + assert ( + TaskState.TASK_STATE_WORKING, + TaskState.TASK_STATE_WORKING, + ) in transitions + assert transitions[-1] == ( + TaskState.TASK_STATE_WORKING, + TaskState.TASK_STATE_COMPLETED, + ) + assert all(i.task_id == result.id for i in infos) + assert all(i.context is context for i in infos) + + names = recorder.names() + assert names[0] == 'before_agent_run' + assert names.index('before_execution') < names.index('on_execution_event') + + +@pytest.mark.asyncio +async def test_streaming_reports_each_event_and_a_final_call(): + recorder = Recorder() + handler = make_handler(recorder) + + streamed = [ + event + async for event in handler.on_message_send_stream( + send_request(), ServerCallContext() + ) + ] + + responses = recorder.hooks('after_agent_run') + assert [r.payload for r in responses[:-1]] == streamed + assert all(r.is_stream_event and not r.is_final for r in responses[:-1]) + final = responses[-1] + assert final.method == 'SendStreamingMessage' + assert final.is_final + assert not final.is_stream_event + assert final.payload is None + assert final.error is None + + +@pytest.mark.asyncio +async def test_stream_closed_early_still_gets_final_call(): + recorder = Recorder() + handler = make_handler(recorder) + + stream = handler.on_message_send_stream(send_request(), ServerCallContext()) + await anext(stream) + await stream.aclose() + + final = recorder.hooks('after_agent_run')[-1] + assert final.is_final + assert final.error is None + + +@pytest.mark.asyncio +async def test_handler_error_reaches_after_agent_run(): + recorder = Recorder() + handler = make_handler(recorder) + + with pytest.raises(TaskNotFoundError): + await handler.on_get_task( + GetTaskRequest(id='missing'), ServerCallContext() + ) + + (response,) = recorder.hooks('after_agent_run') + assert response.method == 'GetTask' + assert isinstance(response.error, TaskNotFoundError) + assert response.is_final + + +@pytest.mark.asyncio +async def test_stream_error_reaches_final_call(): + recorder = Recorder() + handler = make_handler(recorder) + + with pytest.raises(TaskNotFoundError): + async for _ in handler.on_message_send_stream( + SendMessageRequest( + message=Message( + role=Role.ROLE_USER, + message_id='m-2', + task_id='missing', + parts=[Part(text='x')], + ) + ), + ServerCallContext(), + ): + pass + + final = recorder.hooks('after_agent_run')[-1] + assert isinstance(final.error, TaskNotFoundError) + + +@pytest.mark.asyncio +async def test_before_agent_run_can_reject_the_request(): + class Deny(plugins.ServerPlugin): + async def before_agent_run(self, request): + raise PermissionError('denied') + + executed = False + + class Spy(TaskAgent): + async def execute(self, context, event_queue): + nonlocal executed + executed = True + + handler = DefaultRequestHandlerV2( + agent_executor=Spy(), + task_store=InMemoryTaskStore(), + agent_card=CARD, + plugins=plugins.sequential(Deny()), + ) + + with pytest.raises(PermissionError): + await handler.on_message_send(send_request(), ServerCallContext()) + with pytest.raises(PermissionError): + async for _ in handler.on_message_send_stream( + send_request(), ServerCallContext() + ): + pass + assert not executed + + +@pytest.mark.asyncio +async def test_message_only_agent_reports_event_but_no_task_write(): + recorder = Recorder() + handler = make_handler(recorder, executor=MessageAgent()) + + result = await handler.on_message_send(send_request(), ServerCallContext()) + await wait_for(lambda: 'after_execution' in recorder.names()) + + assert isinstance(result, Message) + (event,) = recorder.hooks('on_execution_event') + assert isinstance(event, Message) + assert recorder.hooks('on_task_event') == [] + + +@pytest.mark.asyncio +async def test_before_execution_rejection_fails_the_task(): + class Reject(plugins.ExecutionPlugin): + async def before_execution(self, request_context): + raise RuntimeError('not now') + + recorder = Recorder() + handler = DefaultRequestHandlerV2( + agent_executor=TaskAgent(), + task_store=InMemoryTaskStore(), + agent_card=CARD, + plugins=plugins.sequential(recorder, Reject()), + ) + + with pytest.raises(RuntimeError, match='not now'): + await handler.on_message_send(send_request(), ServerCallContext()) + + # The recorder ran before the rejecting plugin, so it is unwound. + assert recorder.hooks('on_execution_event') == [] + (error,) = recorder.hooks('after_execution') + assert isinstance(error, RuntimeError) + failed = [ + i + for i in recorder.hooks('on_task_event') + if i.to_state == TaskState.TASK_STATE_FAILED + ] + assert failed + + +@pytest.mark.asyncio +async def test_agent_error_reaches_after_execution(): + class Boom(AgentExecutor): + async def execute(self, context, event_queue): + raise ValueError('boom') + + async def cancel(self, context, event_queue): + pass + + recorder = Recorder() + handler = make_handler(recorder, executor=Boom()) + + with pytest.raises(ValueError, match='boom'): + await handler.on_message_send(send_request(), ServerCallContext()) + + (error,) = recorder.hooks('after_execution') + assert isinstance(error, ValueError) + + +@pytest.mark.asyncio +async def test_return_immediately_ends_request_before_execution(): + release = asyncio.Event() + + class SlowAgent(TaskAgent): + async def execute(self, context, event_queue): + task = new_task_from_user_message(context.message) + await event_queue.enqueue_event(task) + await release.wait() + await TaskUpdater(event_queue, task.id, task.context_id).complete() + + recorder = Recorder() + handler = make_handler(recorder, executor=SlowAgent()) + + result = await handler.on_message_send( + send_request( + configuration=SendMessageConfiguration(return_immediately=True) + ), + ServerCallContext(), + ) + + assert result.status.state == TaskState.TASK_STATE_SUBMITTED + assert recorder.hooks('after_agent_run')[0].is_final + assert 'after_execution' not in recorder.names() + + release.set() + # after_execution (producer) and the final task write (consumer) run in + # different tasks, so wait for both. + await wait_for( + lambda: ( + 'after_execution' in recorder.names() + and recorder.hooks('on_task_event')[-1].to_state + == TaskState.TASK_STATE_COMPLETED + ) + ) + + +@pytest.mark.asyncio +async def test_remote_cancel_reports_task_event(): + recorder = Recorder() + store = InMemoryTaskStore() + context = ServerCallContext() + await store.save( + Task( + id='t-remote', + context_id='c-1', + status=TaskStatus(state=TaskState.TASK_STATE_WORKING), + ), + context, + ) + handler = make_handler(recorder, task_store=store) + + result = await handler.on_cancel_task( + CancelTaskRequest(id='t-remote'), context + ) + + assert result.status.state == TaskState.TASK_STATE_CANCELED + (info,) = recorder.hooks('on_task_event') + assert info.from_state == TaskState.TASK_STATE_WORKING + assert info.to_state == TaskState.TASK_STATE_CANCELED + assert info.event.status.state == TaskState.TASK_STATE_CANCELED + + +def push_setup( + responses: dict[str, int], +) -> tuple[ + InMemoryPushNotificationConfigStore, + BasePushNotificationSender, + list[httpx.Request], +]: + received: list[httpx.Request] = [] + + def handle(request: httpx.Request) -> httpx.Response: + received.append(request) + return httpx.Response(responses.get(str(request.url), 200)) + + config_store = InMemoryPushNotificationConfigStore() + sender = BasePushNotificationSender( + httpx.AsyncClient(transport=httpx.MockTransport(handle)), config_store + ) + return config_store, sender, received + + +def push_request(url: str) -> SendMessageRequest: + return send_request( + configuration=SendMessageConfiguration( + task_push_notification_config=TaskPushNotificationConfig( + url=url, token='tok' + ) + ) + ) + + +@pytest.mark.asyncio +async def test_push_hooks_run_for_v2_deliveries(): + config_store, sender, received = push_setup({'http://hook/fail': 500}) + recorder = Recorder() + handler = make_handler( + recorder, push_config_store=config_store, push_sender=sender + ) + + await handler.on_message_send( + push_request('http://hook/ok'), ServerCallContext() + ) + await wait_for(lambda: len(recorder.hooks('on_push_sent')) >= 3) + + assert received + assert all(r.headers['x-plugin'] == 'yes' for r in received) + assert all(r.headers['x-a2a-notification-token'] == 'tok' for r in received) + results = recorder.hooks('on_push_sent') + assert all(r.status_code == 200 and r.error is None for r in results) + assert all(r.duration_s >= 0 for r in results) + assert all(r.request.url == 'http://hook/ok' for r in results) + + recorder.calls.clear() + await handler.on_message_send( + push_request('http://hook/fail'), ServerCallContext() + ) + await wait_for(lambda: len(recorder.hooks('on_push_sent')) >= 1) + failed = recorder.hooks('on_push_sent')[0] + assert failed.status_code == 500 + assert isinstance(failed.error, httpx.HTTPStatusError) + + +@pytest.mark.asyncio +async def test_push_skip_and_rejection(): + class SkipAll(PushPlugin): + async def before_push_send(self, request): + request.skip = True + + class Fail(PushPlugin): + async def before_push_send(self, request): + raise RuntimeError('no') + + seen: list[PushResult] = [] + + class Seen(PushPlugin): + async def on_push_sent(self, result): + seen.append(result) + + for plugin, check in ( + (SkipAll(), lambda r: r.skipped and r.error is None), + (Fail(), lambda r: isinstance(r.error, RuntimeError)), + ): + seen.clear() + config_store, sender, received = push_setup({}) + handler = DefaultRequestHandlerV2( + agent_executor=TaskAgent(), + task_store=InMemoryTaskStore(), + agent_card=CARD, + push_config_store=config_store, + push_sender=sender, + plugins=plugins.sequential(plugin, Seen()), + ) + await handler.on_message_send( + push_request('http://hook/x'), ServerCallContext() + ) + await wait_for(lambda: len(seen) >= 1) + assert received == [] + assert all(check(r) for r in seen) + + +@pytest.mark.asyncio +async def test_push_url_validator_rejection_runs_no_hooks(): + async def reject(url: str) -> bool: + return False + + received: list[httpx.Request] = [] + config_store = InMemoryPushNotificationConfigStore() + sender = BasePushNotificationSender( + httpx.AsyncClient( + transport=httpx.MockTransport( + lambda r: received.append(r) or httpx.Response(200) + ) + ), + config_store, + push_url_validator=reject, + ) + recorder = Recorder() + handler = make_handler( + recorder, push_config_store=config_store, push_sender=sender + ) + + await handler.on_message_send( + push_request('http://hook/x'), ServerCallContext() + ) + await wait_for(lambda: 'after_execution' in recorder.names()) + await asyncio.sleep(0.05) + + assert received == [] + assert recorder.hooks('before_push_send') == [] + + +@pytest.mark.asyncio +async def test_legacy_handler_runs_no_plugin_hooks(): + recorder = Recorder() + plugins.set_global(plugins.sequential(recorder)) + config_store, sender, received = push_setup({}) + handler = LegacyRequestHandler( + agent_executor=TaskAgent(), + task_store=InMemoryTaskStore(), + agent_card=CARD, + push_config_store=config_store, + push_sender=sender, + ) + + result = await handler.on_message_send( + push_request('http://hook/legacy'), ServerCallContext() + ) + await asyncio.sleep(0.1) + + assert result.status.state == TaskState.TASK_STATE_COMPLETED + assert received + assert 'x-plugin' not in received[0].headers + assert recorder.calls == [] diff --git a/tests/server/request_handlers/test_grpc_handler.py b/tests/server/request_handlers/test_grpc_handler.py index fa504fc05..8b6ff2713 100644 --- a/tests/server/request_handlers/test_grpc_handler.py +++ b/tests/server/request_handlers/test_grpc_handler.py @@ -746,3 +746,28 @@ async def mock_stream(*args, **kwargs): server_context = call_args[0][1] assert isinstance(server_context, ServerCallContext) assert server_context.tenant == '' + + +@pytest.mark.asyncio +async def test_call_context_records_transport_and_text_metadata( + grpc_handler: GrpcHandler, + mock_request_handler: AsyncMock, + mock_grpc_context: AsyncMock, +) -> None: + mock_grpc_context.invocation_metadata.return_value = grpc.aio.Metadata( + ('traceparent', '00-abc-def-01'), + ('X-Tenant', 'acme'), + ('trace-bin', b'\x00\x01'), + ) + mock_request_handler.on_get_task.return_value = types.Task(id='t-1') + + await grpc_handler.GetTask( + a2a_pb2.GetTaskRequest(id='t-1'), mock_grpc_context + ) + + call_context = mock_request_handler.on_get_task.call_args[0][1] + assert call_context.state['a2a.transport'] == 'GRPC' + assert call_context.state['headers'] == { + 'traceparent': '00-abc-def-01', + 'x-tenant': 'acme', + } diff --git a/tests/server/tasks/test_task_manager_plugins.py b/tests/server/tasks/test_task_manager_plugins.py new file mode 100644 index 000000000..e7ece4e57 --- /dev/null +++ b/tests/server/tasks/test_task_manager_plugins.py @@ -0,0 +1,163 @@ +from unittest.mock import AsyncMock + +import pytest + +from a2a import plugins +from a2a.plugins import TaskEventInfo, TaskObserverPlugin +from a2a.server.cluster.task_store import ConcurrentTaskModificationError +from a2a.server.context import ServerCallContext +from a2a.server.tasks import InMemoryTaskStore, TaskManager +from a2a.types.a2a_pb2 import ( + Artifact, + Part, + Task, + TaskArtifactUpdateEvent, + TaskState, + TaskStatus, + TaskStatusUpdateEvent, +) + + +class Collect(TaskObserverPlugin): + def __init__(self) -> None: + self.infos: list[TaskEventInfo] = [] + + async def on_task_event(self, info: TaskEventInfo) -> None: + self.infos.append(info) + + def transitions(self) -> list[tuple]: + return [(i.from_state, i.to_state) for i in self.infos] + + +def status(task_id: str, state: TaskState) -> TaskStatusUpdateEvent: + return TaskStatusUpdateEvent( + task_id=task_id, context_id='c-1', status=TaskStatus(state=state) + ) + + +def manager(store, collect, task_id='t-1') -> TaskManager: + return TaskManager( + task_store=store, + context=ServerCallContext(), + task_id=task_id, + context_id='c-1', + initial_message=None, + plugins=plugins.sequential(collect), + ) + + +@pytest.mark.asyncio +async def test_transitions_from_creation_to_completion(): + collect = Collect() + tm = manager(InMemoryTaskStore(), collect) + + await tm.save_task_event(status('t-1', TaskState.TASK_STATE_WORKING)) + await tm.save_task_event( + TaskArtifactUpdateEvent( + task_id='t-1', + context_id='c-1', + artifact=Artifact(artifact_id='a-1', parts=[Part(text='x')]), + ) + ) + await tm.save_task_event(status('t-1', TaskState.TASK_STATE_COMPLETED)) + + W, S, C = ( + TaskState.TASK_STATE_WORKING, + TaskState.TASK_STATE_SUBMITTED, + TaskState.TASK_STATE_COMPLETED, + ) + assert collect.transitions() == [(None, S), (S, W), (W, W), (W, C)] + # The creation write has no event; the others carry the agent event. + assert collect.infos[0].event is None + assert isinstance(collect.infos[1].event, TaskStatusUpdateEvent) + assert isinstance(collect.infos[2].event, TaskArtifactUpdateEvent) + assert collect.infos[-1].task.status.state == C + + +@pytest.mark.asyncio +async def test_from_state_comes_from_existing_task_in_store(): + store = InMemoryTaskStore() + context = ServerCallContext() + await store.save( + Task( + id='t-1', + context_id='c-1', + status=TaskStatus(state=TaskState.TASK_STATE_INPUT_REQUIRED), + ), + context, + ) + collect = Collect() + tm = manager(store, collect) + + await tm.save_task_event(status('t-1', TaskState.TASK_STATE_WORKING)) + + assert collect.transitions() == [ + (TaskState.TASK_STATE_INPUT_REQUIRED, TaskState.TASK_STATE_WORKING) + ] + + +@pytest.mark.asyncio +async def test_get_task_and_invalidate_refresh_persisted_state(): + store = InMemoryTaskStore() + context = ServerCallContext() + collect = Collect() + tm = manager(store, collect) + assert await tm.get_task() is None + + await tm.save_task_event( + Task( + id='t-1', + context_id='c-1', + status=TaskStatus(state=TaskState.TASK_STATE_SUBMITTED), + ) + ) + # Another writer advances the task behind the manager's back. + await store.save( + Task( + id='t-1', + context_id='c-1', + status=TaskStatus(state=TaskState.TASK_STATE_WORKING), + ), + context, + ) + tm.invalidate() + assert (await tm.get_task()).status.state == TaskState.TASK_STATE_WORKING + await tm.save_task_event(status('t-1', TaskState.TASK_STATE_COMPLETED)) + + assert collect.transitions() == [ + (None, TaskState.TASK_STATE_SUBMITTED), + (TaskState.TASK_STATE_WORKING, TaskState.TASK_STATE_COMPLETED), + ] + + +@pytest.mark.asyncio +async def test_failed_write_reports_nothing(): + collect = Collect() + store = InMemoryTaskStore() + tm = manager(store, collect) + tm._versioned_store.save = AsyncMock( + side_effect=ConcurrentTaskModificationError('stale') + ) + + with pytest.raises(ConcurrentTaskModificationError): + await tm.save_task_event(status('t-1', TaskState.TASK_STATE_WORKING)) + + assert collect.infos == [] + + +@pytest.mark.asyncio +async def test_no_plugins_or_no_task_hooks_reports_nothing(): + class OnlyServer(plugins.ServerPlugin): + pass + + for manager_arg in (None, plugins.sequential(OnlyServer())): + tm = TaskManager( + task_store=InMemoryTaskStore(), + context=ServerCallContext(), + task_id='t-1', + context_id='c-1', + initial_message=None, + plugins=manager_arg, + ) + assert tm._plugins is None + await tm.save_task_event(status('t-1', TaskState.TASK_STATE_WORKING)) From 114e7817e05e88ec06826c4a96f8fe1c75c40f99 Mon Sep 17 00:00:00 2001 From: Mykyta Netipa Date: Sat, 17 Oct 2026 15:17:52 +0000 Subject: [PATCH 03/11] feat(client): run plugin hooks in BaseClient Every client call now runs before_agent_call before the interceptors and after_agent_call after them, with one call per stream event, a final call when the stream ends, and the error when the transport fails. The hooks get the RPC name used on the server side and the protocol binding and URL ClientFactory selected, so a plugin can tag calls and add headers such as trace context. send_message and subscribe now close the inner stream with aclosing() so an early break by the caller releases the transport stream and runs the final hook right away instead of at garbage collection. --- src/a2a/client/base_client.py | 186 +++++++++++-- src/a2a/client/client_factory.py | 2 + tests/client/test_base_client_plugins.py | 318 +++++++++++++++++++++++ 3 files changed, 485 insertions(+), 21 deletions(-) create mode 100644 tests/client/test_base_client_plugins.py diff --git a/src/a2a/client/base_client.py b/src/a2a/client/base_client.py index 763f23fb5..56b5737ed 100644 --- a/src/a2a/client/base_client.py +++ b/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, @@ -31,21 +34,61 @@ ) +# Client method names mapped to the A2A RPC names reported to plugins. +_RPC_NAMES = { + 'send_message': 'SendMessage', + 'send_message_streaming': 'SendStreamingMessage', + 'get_task': 'GetTask', + 'list_tasks': 'ListTasks', + 'cancel_task': 'CancelTask', + 'create_task_push_notification_config': 'CreateTaskPushNotificationConfig', + 'get_task_push_notification_config': 'GetTaskPushNotificationConfig', + 'list_task_push_notification_configs': 'ListTaskPushNotificationConfigs', + 'delete_task_push_notification_config': 'DeleteTaskPushNotificationConfig', + 'subscribe': 'SubscribeToTask', + 'get_extended_agent_card': 'GetExtendedAgentCard', +} + + class BaseClient(Client): - """Base implementation of the A2A client, containing transport-independent logic.""" + """Base implementation of the A2A client, containing transport-independent logic. - def __init__( + Client plugin hooks run around every call: ``before_agent_call`` before + the interceptors, ``after_agent_call`` after them. + """ + + def __init__( # noqa: PLR0913 self, card: AgentCard, config: ClientConfig, transport: ClientTransport, interceptors: list[ClientCallInterceptor], + *, + plugins: PluginManager | None = None, + protocol_binding: str = '', + url: str = '', ): + """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. + protocol_binding: Protocol binding of ``transport``, reported to + plugins. + url: URL of the selected agent interface, reported to plugins. + """ 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() + self._protocol_binding = protocol_binding + self._url = url async def send_message( self, @@ -90,15 +133,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 +372,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 +423,18 @@ 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=_RPC_NAMES.get(method, method), + payload=input_data, + agent_card=self._card, + context=context if context is not None else ClientCallContext(), + transport=self._protocol_binding, + url=self._url, + ) + async def _execute_with_interceptors( self, input_data: Any, @@ -380,6 +443,37 @@ async def _execute_with_interceptors( transport_call: Callable[ [Any, ClientCallContext | None], Awaitable[Any] ], + ) -> Any: + 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_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 +516,58 @@ async def _execute_stream_with_interceptors( transport_call: Callable[ [Any, ClientCallContext | None], AsyncIterator[StreamResponse] ], - ) -> AsyncIterator[StreamResponse]: + ) -> AsyncGenerator[StreamResponse, None]: + 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_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, diff --git a/src/a2a/client/client_factory.py b/src/a2a/client/client_factory.py index e43207ca7..424402bcc 100644 --- a/src/a2a/client/client_factory.py +++ b/src/a2a/client/client_factory.py @@ -369,6 +369,8 @@ def create( self._config, transport, interceptors or [], + protocol_binding=transport_protocol, + url=selected_interface.url, ) diff --git a/tests/client/test_base_client_plugins.py b/tests/client/test_base_client_plugins.py new file mode 100644 index 000000000..5f184995b --- /dev/null +++ b/tests/client/test_base_client_plugins.py @@ -0,0 +1,318 @@ +from typing import Any +from unittest.mock import AsyncMock + +import httpx +import pytest + +from a2a import plugins +from a2a.client.base_client import BaseClient +from a2a.client.client import ClientCallContext, ClientConfig +from a2a.client.client_factory import ClientFactory +from a2a.client.interceptors import AfterArgs, BeforeArgs, ClientCallInterceptor +from a2a.client.transports.base import ClientTransport +from a2a.plugins import ClientRequest, ClientResponse +from a2a.types.a2a_pb2 import ( + AgentCapabilities, + AgentCard, + AgentInterface, + GetTaskRequest, + Message, + Part, + Role, + SendMessageRequest, + SendMessageResponse, + StreamResponse, + Task, + TaskState, + TaskStatus, + TaskStatusUpdateEvent, +) + + +class Recorder(plugins.ClientPlugin): + def __init__(self) -> None: + self.calls: list[tuple[str, Any]] = [] + + def hooks(self, name: str) -> list[Any]: + return [arg for hook, arg in self.calls if hook == name] + + async def before_agent_call(self, request: ClientRequest) -> None: + params = dict(request.context.service_parameters or {}) + params['x-trace'] = 'abc' + request.context.service_parameters = params + self.calls.append(('before_agent_call', request)) + + async def after_agent_call(self, response: ClientResponse) -> None: + self.calls.append(('after_agent_call', response)) + + +def card(streaming: bool = True) -> AgentCard: + return AgentCard( + name='Remote', + version='1.0', + capabilities=AgentCapabilities(streaming=streaming), + supported_interfaces=[ + AgentInterface(url='http://remote:9000', protocol_binding='JSONRPC') + ], + ) + + +def client( + recorder: Recorder | None, + transport: AsyncMock, + interceptors: list[ClientCallInterceptor] | None = None, + streaming: bool = True, +) -> BaseClient: + return BaseClient( + card(streaming), + ClientConfig(streaming=streaming), + transport, + interceptors or [], + plugins=plugins.sequential(recorder) if recorder else None, + protocol_binding='JSONRPC', + url='http://remote:9000', + ) + + +def send_request() -> SendMessageRequest: + return SendMessageRequest( + message=Message( + role=Role.ROLE_USER, message_id='m-1', parts=[Part(text='hi')] + ) + ) + + +def status_event(state: TaskState) -> StreamResponse: + return StreamResponse( + status_update=TaskStatusUpdateEvent( + task_id='t-1', status=TaskStatus(state=state) + ) + ) + + +@pytest.fixture(autouse=True) +def restore_global(): + previous = plugins.get_global() + yield + plugins.set_global(previous) + + +@pytest.mark.asyncio +async def test_unary_call_runs_hooks_with_transport_details(): + recorder = Recorder() + transport = AsyncMock(spec=ClientTransport) + transport.get_task.return_value = Task(id='t-1') + + result = await client(recorder, transport).get_task( + GetTaskRequest(id='t-1') + ) + + assert result.id == 't-1' + (request,) = recorder.hooks('before_agent_call') + assert request.method == 'GetTask' + assert request.transport == 'JSONRPC' + assert request.url == 'http://remote:9000' + assert request.agent_card.name == 'Remote' + (response,) = recorder.hooks('after_agent_call') + assert response.payload is result + assert response.is_final + assert response.context is request.context + # The header added by the plugin reaches the transport. + sent_context = transport.get_task.call_args.kwargs['context'] + assert sent_context.service_parameters['x-trace'] == 'abc' + + +@pytest.mark.asyncio +async def test_non_streaming_send_message_is_a_unary_call(): + recorder = Recorder() + transport = AsyncMock(spec=ClientTransport) + transport.send_message.return_value = SendMessageResponse( + task=Task(id='t-1') + ) + + events = [ + e + async for e in client( + recorder, transport, streaming=False + ).send_message(send_request()) + ] + + assert len(events) == 1 + assert recorder.hooks('before_agent_call')[0].method == 'SendMessage' + (response,) = recorder.hooks('after_agent_call') + assert response.is_final + + +@pytest.mark.asyncio +async def test_streaming_reports_events_and_final_call(): + recorder = Recorder() + transport = AsyncMock(spec=ClientTransport) + streamed = [ + status_event(TaskState.TASK_STATE_WORKING), + status_event(TaskState.TASK_STATE_COMPLETED), + ] + + async def stream(request, *, context=None): + for e in streamed: + yield e + + transport.send_message_streaming = stream + + events = [ + e + async for e in client(recorder, transport).send_message(send_request()) + ] + + assert events == streamed + (request,) = recorder.hooks('before_agent_call') + assert request.method == 'SendStreamingMessage' + responses = recorder.hooks('after_agent_call') + assert [r.payload for r in responses[:-1]] == streamed + assert all(r.is_stream_event and not r.is_final for r in responses[:-1]) + assert responses[-1].is_final + assert responses[-1].payload is None + assert responses[-1].error is None + + +@pytest.mark.asyncio +async def test_stream_stopped_early_gets_final_call(): + recorder = Recorder() + transport = AsyncMock(spec=ClientTransport) + + async def stream(request, *, context=None): + for _ in range(10): + yield status_event(TaskState.TASK_STATE_WORKING) + + transport.send_message_streaming = stream + events = client(recorder, transport).send_message(send_request()) + await anext(events) + await events.aclose() + + assert recorder.hooks('after_agent_call')[-1].is_final + + +@pytest.mark.asyncio +async def test_transport_errors_reach_after_agent_call(): + recorder = Recorder() + transport = AsyncMock(spec=ClientTransport) + transport.get_task.side_effect = ConnectionError('down') + + with pytest.raises(ConnectionError): + await client(recorder, transport).get_task(GetTaskRequest(id='t-1')) + + (response,) = recorder.hooks('after_agent_call') + assert isinstance(response.error, ConnectionError) + + async def failing(request, *, context=None): + yield status_event(TaskState.TASK_STATE_WORKING) + raise ConnectionError('dropped') + + transport.send_message_streaming = failing + recorder.calls.clear() + with pytest.raises(ConnectionError): + async for _ in client(recorder, transport).send_message(send_request()): + pass + assert isinstance( + recorder.hooks('after_agent_call')[-1].error, ConnectionError + ) + + +@pytest.mark.asyncio +async def test_plugins_wrap_interceptors_including_early_return(): + order: list[str] = [] + + class Order(plugins.ClientPlugin): + async def before_agent_call(self, request): + order.append('plugin.before') + + async def after_agent_call(self, response): + order.append('plugin.after') + + class Cached(ClientCallInterceptor): + async def before(self, args: BeforeArgs) -> None: + order.append('interceptor.before') + args.early_return = Task(id='cached') + + async def after(self, args: AfterArgs) -> None: + order.append('interceptor.after') + + transport = AsyncMock(spec=ClientTransport) + c = BaseClient( + card(), + ClientConfig(), + transport, + [Cached()], + plugins=plugins.sequential(Order()), + ) + + result = await c.get_task(GetTaskRequest(id='t-1')) + + assert result.id == 'cached' + transport.get_task.assert_not_called() + assert order == [ + 'plugin.before', + 'interceptor.before', + 'interceptor.after', + 'plugin.after', + ] + + +@pytest.mark.asyncio +async def test_before_agent_call_can_cancel_the_call(): + class Deny(plugins.ClientPlugin): + async def before_agent_call(self, request): + raise PermissionError('no') + + transport = AsyncMock(spec=ClientTransport) + c = BaseClient( + card(), + ClientConfig(), + transport, + [], + plugins=plugins.sequential(Deny()), + ) + + with pytest.raises(PermissionError): + await c.get_task(GetTaskRequest(id='t-1')) + with pytest.raises(PermissionError): + async for _ in c.send_message(send_request()): + pass + transport.get_task.assert_not_called() + + +@pytest.mark.asyncio +async def test_existing_context_is_passed_through(): + recorder = Recorder() + transport = AsyncMock(spec=ClientTransport) + transport.get_task.return_value = Task(id='t-1') + context = ClientCallContext(state={'k': 'v'}) + + await client(recorder, transport).get_task( + GetTaskRequest(id='t-1'), context=context + ) + + assert recorder.hooks('before_agent_call')[0].context is context + + +@pytest.mark.asyncio +async def test_no_plugins_keeps_context_none(): + transport = AsyncMock(spec=ClientTransport) + transport.get_task.return_value = Task(id='t-1') + + await client(None, transport).get_task(GetTaskRequest(id='t-1')) + + assert transport.get_task.call_args.kwargs['context'] is None + + +def test_factory_passes_selected_transport_and_global_plugins(): + manager = plugins.sequential(Recorder()) + plugins.set_global(manager) + + c = ClientFactory(ClientConfig(httpx_client=httpx.AsyncClient())).create( + card() + ) + + assert isinstance(c, BaseClient) + assert c._plugins is manager + assert c._protocol_binding == 'JSONRPC' + assert c._url == 'http://remote:9000' From 8a22fbfd58b5093aa94004bbf1fb4fd4fd4774f9 Mon Sep 17 00:00:00 2001 From: Mykyta Netipa Date: Sat, 17 Oct 2026 15:17:52 +0000 Subject: [PATCH 04/11] test: cover plugins over JSON-RPC, HTTP+JSON, gRPC and v0.3 Runs client and server hooks against in-process servers on every transport, checks that a header added by a client plugin reaches the server, that v0.3 JSON-RPC requests reach the V2 hooks, and that state set in before_execution propagates to calls the agent makes to a second agent. --- tests/integration/test_plugins.py | 401 ++++++++++++++++++++++++++++++ 1 file changed, 401 insertions(+) create mode 100644 tests/integration/test_plugins.py diff --git a/tests/integration/test_plugins.py b/tests/integration/test_plugins.py new file mode 100644 index 000000000..86f9cf556 --- /dev/null +++ b/tests/integration/test_plugins.py @@ -0,0 +1,401 @@ +"""Plugins over real transports: JSON-RPC, HTTP+JSON, gRPC and v0.3 JSON-RPC.""" + +import asyncio +import contextvars + +from collections.abc import AsyncGenerator, Callable +from typing import Any, NamedTuple + +import grpc +import httpx +import pytest +import pytest_asyncio + +from a2a import plugins +from a2a.client import ClientConfig, ClientFactory +from a2a.client.base_client import BaseClient +from a2a.helpers.proto_helpers import new_task_from_user_message +from a2a.plugins import ( + ClientRequest, + ClientResponse, + CompletePlugin, + ExecutionPlugin, + ServerRequest, + ServerResponse, + TaskEventInfo, + TaskObserverPlugin, +) +from a2a.server.agent_execution import AgentExecutor, RequestContext +from a2a.server.events import EventQueue +from a2a.server.request_handlers import DefaultRequestHandlerV2, GrpcHandler +from a2a.server.routes import ( + create_agent_card_routes, + create_jsonrpc_routes, + create_rest_routes, +) +from a2a.server.tasks import InMemoryTaskStore, TaskUpdater +from a2a.types import a2a_pb2_grpc +from a2a.types.a2a_pb2 import ( + AgentCapabilities, + AgentCard, + AgentInterface, + GetTaskRequest, + Message, + Part, + Role, + SendMessageRequest, + TaskState, +) +from a2a.utils import TransportProtocol +from starlette.applications import Starlette + + +class Recorder(CompletePlugin, ExecutionPlugin, TaskObserverPlugin): + """Records hooks and adds an ``x-trace`` header to outgoing calls.""" + + name = 'recorder' + + def __init__(self) -> None: + self.calls: list[tuple[str, Any]] = [] + + def hooks(self, name: str) -> list[Any]: + return [arg for hook, arg in self.calls if hook == name] + + async def before_agent_run(self, request: ServerRequest) -> None: + self.calls.append(('before_agent_run', request)) + + async def after_agent_run(self, response: ServerResponse) -> None: + self.calls.append(('after_agent_run', response)) + + async def before_agent_call(self, request: ClientRequest) -> None: + params = dict(request.context.service_parameters or {}) + params['x-trace'] = 'trace-123' + request.context.service_parameters = params + self.calls.append(('before_agent_call', request)) + + async def after_agent_call(self, response: ClientResponse) -> None: + self.calls.append(('after_agent_call', response)) + + async def before_execution(self, request_context: RequestContext) -> None: + self.calls.append(('before_execution', request_context)) + + async def on_execution_event(self, request_context, event) -> None: + self.calls.append(('on_execution_event', event)) + + async def after_execution(self, request_context, error) -> None: + self.calls.append(('after_execution', error)) + + async def on_task_event(self, info: TaskEventInfo) -> None: + self.calls.append(('on_task_event', info)) + + +class TaskAgent(AgentExecutor): + async def execute(self, context: RequestContext, event_queue: EventQueue): + task = context.current_task or new_task_from_user_message( + context.message + ) + if not context.current_task: + await event_queue.enqueue_event(task) + updater = TaskUpdater(event_queue, task.id, task.context_id) + await updater.start_work() + await updater.add_artifact(parts=[Part(text='done')], name='out') + await updater.complete() + + async def cancel(self, context: RequestContext, event_queue: EventQueue): + pass + + +def agent_card(grpc_url: str = 'localhost:0') -> AgentCard: + return AgentCard( + name='Plugin Agent', + description='Integration test agent', + version='1.0', + capabilities=AgentCapabilities(streaming=True), + default_input_modes=['text/plain'], + default_output_modes=['text/plain'], + supported_interfaces=[ + AgentInterface( + protocol_binding=TransportProtocol.HTTP_JSON, + url='http://testserver', + ), + AgentInterface( + protocol_binding=TransportProtocol.JSONRPC, + url='http://testserver', + ), + AgentInterface( + protocol_binding=TransportProtocol.GRPC, url=grpc_url + ), + ], + ) + + +def send_request(text: str = 'hello') -> SendMessageRequest: + return SendMessageRequest( + message=Message( + role=Role.ROLE_USER, message_id='m-1', parts=[Part(text=text)] + ) + ) + + +async def wait_for(predicate: Callable[[], bool], timeout: float = 5) -> None: + async def poll() -> None: + while not predicate(): + await asyncio.sleep(0.01) + + await asyncio.wait_for(poll(), timeout) + + +class Setup(NamedTuple): + client: BaseClient + server: Recorder + caller: Recorder + transport: str + + +def http_app(handler: DefaultRequestHandlerV2, card: AgentCard) -> Starlette: + return Starlette( + routes=[ + *create_agent_card_routes(agent_card=card, card_url='/card'), + *create_jsonrpc_routes( + request_handler=handler, rpc_url='/', enable_v0_3_compat=True + ), + *create_rest_routes(request_handler=handler), + ] + ) + + +@pytest.fixture(autouse=True) +def restore_global(): + previous = plugins.get_global() + yield + plugins.set_global(previous) + + +@pytest_asyncio.fixture( + params=[ + pytest.param(TransportProtocol.JSONRPC, id='JSON-RPC'), + pytest.param(TransportProtocol.HTTP_JSON, id='HTTP+JSON'), + pytest.param(TransportProtocol.GRPC, id='gRPC'), + ] +) +async def setup(request) -> AsyncGenerator[Setup, None]: + transport = request.param + server_recorder, caller_recorder = Recorder(), Recorder() + server = None + card = agent_card() + handler = DefaultRequestHandlerV2( + agent_executor=TaskAgent(), + task_store=InMemoryTaskStore(), + agent_card=card, + plugins=plugins.sequential(server_recorder), + ) + # Clients pick up the global manager when created. + plugins.set_global(plugins.sequential(caller_recorder)) + + if transport == TransportProtocol.GRPC: + server = grpc.aio.server() + port = server.add_insecure_port('[::]:0') + card = agent_card(f'localhost:{port}') + a2a_pb2_grpc.add_A2AServiceServicer_to_server( + GrpcHandler(handler), server + ) + await server.start() + config = ClientConfig( + grpc_channel_factory=grpc.aio.insecure_channel, + supported_protocol_bindings=[transport], + ) + else: + config = ClientConfig( + httpx_client=httpx.AsyncClient( + transport=httpx.ASGITransport(app=http_app(handler, card)), + base_url='http://testserver', + ), + supported_protocol_bindings=[transport], + ) + + client = ClientFactory(config).create(card) + yield Setup(client, server_recorder, caller_recorder, transport) + + await client.close() + if server is not None: + await server.stop(0) + + +@pytest.mark.asyncio +async def test_streaming_message_runs_hooks_on_both_sides(setup: Setup): + events = [e async for e in setup.client.send_message(send_request())] + await wait_for(lambda: bool(setup.server.hooks('after_execution'))) + + assert events[-1].status_update.status.state == ( + TaskState.TASK_STATE_COMPLETED + ) + + # Client side. + (call,) = setup.caller.hooks('before_agent_call') + assert call.method == 'SendStreamingMessage' + assert call.transport == setup.transport + assert call.url + client_responses = setup.caller.hooks('after_agent_call') + assert len(client_responses) == len(events) + 1 + assert client_responses[-1].is_final + + # Server side: transport recorded, client header received. + (run,) = setup.server.hooks('before_agent_run') + assert run.method == 'SendStreamingMessage' + assert run.context.state['a2a.transport'] == setup.transport + assert run.context.state['headers']['x-trace'] == 'trace-123' + assert run.agent_card.name == 'Plugin Agent' + server_responses = setup.server.hooks('after_agent_run') + assert server_responses[-1].is_final + assert sum(r.is_stream_event for r in server_responses) == len(events) + + # Execution and task hooks. + assert len(setup.server.hooks('on_execution_event')) == 4 + transitions = [ + (i.from_state, i.to_state) for i in setup.server.hooks('on_task_event') + ] + assert transitions[0] == (None, TaskState.TASK_STATE_SUBMITTED) + assert transitions[-1][1] == TaskState.TASK_STATE_COMPLETED + + +@pytest.mark.asyncio +async def test_unary_call_runs_hooks_on_both_sides(setup: Setup): + events = [e async for e in setup.client.send_message(send_request())] + task_id = events[0].task.id + setup.server.calls.clear() + setup.caller.calls.clear() + + task = await setup.client.get_task(GetTaskRequest(id=task_id)) + + assert task.id == task_id + assert [r.method for r in setup.caller.hooks('before_agent_call')] == [ + 'GetTask' + ] + (response,) = setup.server.hooks('after_agent_run') + assert response.method == 'GetTask' + assert response.payload.id == task_id + assert response.is_final + + +@pytest.mark.asyncio +async def test_v03_jsonrpc_requests_run_server_hooks(): + recorder = Recorder() + handler = DefaultRequestHandlerV2( + agent_executor=TaskAgent(), + task_store=InMemoryTaskStore(), + agent_card=agent_card(), + plugins=plugins.sequential(recorder), + ) + async with httpx.AsyncClient( + transport=httpx.ASGITransport(app=http_app(handler, agent_card())), + base_url='http://testserver', + ) as http: + response = await http.post( + '/', + json={ + 'jsonrpc': '2.0', + 'id': '1', + 'method': 'message/send', + 'params': { + 'message': { + 'messageId': 'm-03', + 'role': 'user', + 'parts': [{'kind': 'text', 'text': 'hello'}], + } + }, + }, + ) + + assert response.status_code == 200 + assert response.json()['result']['status']['state'] == 'completed' + (run,) = recorder.hooks('before_agent_run') + assert run.method == 'SendMessage' + assert run.context.state['a2a.transport'] == 'JSONRPC' + assert recorder.hooks('after_agent_run')[-1].is_final + + +_tenant: contextvars.ContextVar[str | None] = contextvars.ContextVar( + 'tenant', default=None +) + + +class TenantPlugin(CompletePlugin, ExecutionPlugin): + """Forwards the incoming x-tenant header to outgoing calls.""" + + name = 'tenant' + + async def before_execution(self, request_context: RequestContext) -> None: + headers = request_context.call_context.state.get('headers', {}) + _tenant.set(headers.get('x-tenant')) + + async def before_agent_call(self, request: ClientRequest) -> None: + if tenant := _tenant.get(): + params = dict(request.context.service_parameters or {}) + params['x-tenant'] = tenant + request.context.service_parameters = params + + +@pytest.mark.asyncio +async def test_two_agent_chain_propagates_state_through_execution(): + downstream = Recorder() + downstream_handler = DefaultRequestHandlerV2( + agent_executor=TaskAgent(), + task_store=InMemoryTaskStore(), + agent_card=agent_card(), + plugins=plugins.sequential(downstream), + ) + downstream_http = httpx.AsyncClient( + transport=httpx.ASGITransport( + app=http_app(downstream_handler, agent_card()) + ), + base_url='http://testserver', + ) + tenant_plugin = TenantPlugin() + plugins.set_global(plugins.sequential(tenant_plugin)) + + class Orchestrator(AgentExecutor): + async def execute(self, context, event_queue): + task = new_task_from_user_message(context.message) + await event_queue.enqueue_event(task) + client = ClientFactory( + ClientConfig(httpx_client=downstream_http) + ).create(agent_card()) + async for _ in client.send_message(send_request('forwarded')): + pass + await TaskUpdater(event_queue, task.id, task.context_id).complete() + + async def cancel(self, context, event_queue): + pass + + upstream_handler = DefaultRequestHandlerV2( + agent_executor=Orchestrator(), + task_store=InMemoryTaskStore(), + agent_card=agent_card(), + ) + async with httpx.AsyncClient( + transport=httpx.ASGITransport( + app=http_app(upstream_handler, agent_card()) + ), + base_url='http://testserver', + ) as http: + response = await http.post( + '/', + headers={'x-tenant': 'acme', 'A2A-Version': '1.0'}, + json={ + 'jsonrpc': '2.0', + 'id': '1', + 'method': 'SendMessage', + 'params': { + 'message': { + 'messageId': 'm-up', + 'role': 'ROLE_USER', + 'parts': [{'text': 'plan'}], + } + }, + }, + ) + await downstream_http.aclose() + + assert response.status_code == 200, response.text + (run,) = downstream.hooks('before_agent_run') + assert run.context.state['headers']['x-tenant'] == 'acme' From 07b8ae0c223eb4f5a478baab2d9aa1291acdfb6e Mon Sep 17 00:00:00 2001 From: Mykyta Netipa Date: Sat, 17 Oct 2026 18:15:33 +0000 Subject: [PATCH 05/11] refactor(server): merge plugin and default push delivery paths Replace the separate _post_with_plugins method with a single delivery path in _dispatch_notification. Header building moves to _build_headers and the on_push_sent call to _report, which both paths share. Without a push plugin manager the only added cost per delivery is constructing a PushRequest and reading a monotonic clock; behavior is unchanged. --- .../tasks/base_push_notification_sender.py | 149 ++++++++---------- 1 file changed, 69 insertions(+), 80 deletions(-) diff --git a/src/a2a/server/tasks/base_push_notification_sender.py b/src/a2a/server/tasks/base_push_notification_sender.py index 26c9c0852..b425251c4 100644 --- a/src/a2a/server/tasks/base_push_notification_sender.py +++ b/src/a2a/server/tasks/base_push_notification_sender.py @@ -126,77 +126,35 @@ async def _dispatch_notification( if not accepted: return False plugins = current_push_plugins() - try: - headers: dict[str, str] = {} - if push_info.token: - headers['X-A2A-Notification-Token'] = push_info.token - auth = push_info.authentication - if push_info.HasField('authentication'): - if auth.scheme and auth.credentials: - headers['Authorization'] = ( - f'{auth.scheme} {auth.credentials}' - ) - elif auth.scheme: - logger.warning( - 'Push config %s sets an authentication scheme with no ' - 'credentials; sending no Authorization header for ' - 'task_id=%s', - push_info.id, - task_id, - ) - - if plugins is not None: - return await self._post_with_plugins( - plugins, PushRequest(task_id, url, headers, event) - ) - response = await self._client.post( - url, - json=MessageToDict(to_stream_response(event)), - headers=headers, - ) - response.raise_for_status() - logger.info( - 'Push-notification sent for task_id=%s to URL: %s', task_id, url - ) - except Exception: - logger.exception( - 'Error sending push-notification for task_id=%s to URL: %s.', - task_id, - url, - ) - return False - return True - - async def _post_with_plugins( - self, plugins: PluginManager, request: PushRequest - ) -> bool: - """Delivers one notification with `PushPlugin` hooks. + request = PushRequest( + task_id, url, self._build_headers(push_info, task_id), event + ) - A raising `before_push_send` fails the delivery. The webhook URL - cannot be changed by plugins, since it was validated already. - """ - url = request.url - try: - await plugins.before_push_send(request) - except Exception as e: - logger.exception( - 'Push plugin failed before sending for task_id=%s to URL: %s.', - request.task_id, - request.url, - ) - await plugins.on_push_sent(PushResult(request, 0.0, None, e)) - return False - if request.skip: - logger.debug( - 'Push-notification for task_id=%s to URL: %s skipped by a ' - 'plugin', - request.task_id, - request.url, - ) - await plugins.on_push_sent( - PushResult(request, 0.0, None, None, skipped=True) - ) - return True + if plugins is not None: + # before_push_send may add headers or set request.skip. The + # validated URL is captured above and cannot be redirected. + try: + await plugins.before_push_send(request) + except Exception as e: + logger.exception( + 'Push plugin failed before sending for task_id=%s to ' + 'URL: %s.', + task_id, + url, + ) + await self._report(plugins, PushResult(request, 0.0, None, e)) + return False + if request.skip: + logger.debug( + 'Push-notification for task_id=%s to URL: %s skipped by a ' + 'plugin', + task_id, + url, + ) + await self._report( + plugins, PushResult(request, 0.0, None, None, skipped=True) + ) + return True start = time.monotonic() status_code: int | None = None @@ -204,7 +162,7 @@ async def _post_with_plugins( try: response = await self._client.post( url, - json=MessageToDict(to_stream_response(request.event)), + json=MessageToDict(to_stream_response(event)), headers=request.headers, ) status_code = response.status_code @@ -213,16 +171,47 @@ async def _post_with_plugins( error = e logger.exception( 'Error sending push-notification for task_id=%s to URL: %s.', - request.task_id, - request.url, + task_id, + url, ) - await plugins.on_push_sent( - PushResult(request, time.monotonic() - start, status_code, error) - ) - if error is None: + else: logger.info( - 'Push-notification sent for task_id=%s to URL: %s', - request.task_id, - request.url, + 'Push-notification sent for task_id=%s to URL: %s', task_id, url + ) + if plugins is not None: + await self._report( + plugins, + PushResult( + request, time.monotonic() - start, status_code, error + ), ) return error is None + + def _build_headers( + self, push_info: TaskPushNotificationConfig, task_id: str + ) -> dict[str, str]: + """Builds the delivery headers from the push config.""" + headers: dict[str, str] = {} + if push_info.token: + headers['X-A2A-Notification-Token'] = push_info.token + auth = push_info.authentication + if push_info.HasField('authentication'): + if auth.scheme and auth.credentials: + headers['Authorization'] = f'{auth.scheme} {auth.credentials}' + elif auth.scheme: + logger.warning( + 'Push config %s sets an authentication scheme with no ' + 'credentials; sending no Authorization header for ' + 'task_id=%s', + push_info.id, + task_id, + ) + return headers + + @staticmethod + async def _report(plugins: PluginManager, result: PushResult) -> None: + """Calls on_push_sent, keeping a raising manager from failing delivery.""" + try: + await plugins.on_push_sent(result) + except Exception: + logger.exception('Push plugin failed in on_push_sent') From 7912da8e937c43ab2f3ac4abd9725c67909f6936 Mon Sep 17 00:00:00 2001 From: Mykyta Netipa Date: Sat, 17 Oct 2026 18:32:07 +0000 Subject: [PATCH 06/11] refactor: drop telemetry-only transport metadata from the plugin PR The plugin hooks do not depend on these additions; they were added ahead of a telemetry plugin that is not part of this change and can be added back compatibly when it lands: - ClientRequest.transport/url and the ClientFactory wiring that filled them; - ServerCallContext.state['a2a.transport'] written by the JSON-RPC, REST, gRPC and v0.3 dispatchers; - gRPC invocation metadata copied into state['headers']. Plugins can still read it through state['grpc_context']. The dispatchers, v0.3 adapters, gRPC handler, ClientFactory and ServerCallContext are back to their state on main. --- src/a2a/client/base_client.py | 11 +------- src/a2a/client/client_factory.py | 2 -- src/a2a/compat/v0_3/grpc_handler.py | 3 --- src/a2a/compat/v0_3/jsonrpc_adapter.py | 6 +---- src/a2a/compat/v0_3/rest_adapter.py | 8 ------ src/a2a/plugins/types.py | 8 ++---- src/a2a/server/context.py | 3 --- .../server/request_handlers/grpc_handler.py | 18 +------------ src/a2a/server/routes/jsonrpc_dispatcher.py | 6 +---- src/a2a/server/routes/rest_dispatcher.py | 6 +---- tests/client/test_base_client_plugins.py | 10 ++------ tests/integration/test_plugins.py | 16 +++++++----- .../request_handlers/test_grpc_handler.py | 25 ------------------- 13 files changed, 19 insertions(+), 103 deletions(-) diff --git a/src/a2a/client/base_client.py b/src/a2a/client/base_client.py index 56b5737ed..4637416b9 100644 --- a/src/a2a/client/base_client.py +++ b/src/a2a/client/base_client.py @@ -57,7 +57,7 @@ class BaseClient(Client): the interceptors, ``after_agent_call`` after them. """ - def __init__( # noqa: PLR0913 + def __init__( self, card: AgentCard, config: ClientConfig, @@ -65,8 +65,6 @@ def __init__( # noqa: PLR0913 interceptors: list[ClientCallInterceptor], *, plugins: PluginManager | None = None, - protocol_binding: str = '', - url: str = '', ): """Initializes the client. @@ -77,9 +75,6 @@ def __init__( # noqa: PLR0913 interceptors: Interceptors applied to every call. plugins: Plugin manager whose client hooks run for every call. Defaults to the global manager at construction time. - protocol_binding: Protocol binding of ``transport``, reported to - plugins. - url: URL of the selected agent interface, reported to plugins. """ super().__init__(interceptors) self._card = card @@ -87,8 +82,6 @@ def __init__( # noqa: PLR0913 self._transport = transport self._interceptors = interceptors self._plugins = plugins if plugins is not None else get_global() - self._protocol_binding = protocol_binding - self._url = url async def send_message( self, @@ -431,8 +424,6 @@ def _client_request( payload=input_data, agent_card=self._card, context=context if context is not None else ClientCallContext(), - transport=self._protocol_binding, - url=self._url, ) async def _execute_with_interceptors( diff --git a/src/a2a/client/client_factory.py b/src/a2a/client/client_factory.py index 424402bcc..e43207ca7 100644 --- a/src/a2a/client/client_factory.py +++ b/src/a2a/client/client_factory.py @@ -369,8 +369,6 @@ def create( self._config, transport, interceptors or [], - protocol_binding=transport_protocol, - url=selected_interface.url, ) diff --git a/src/a2a/compat/v0_3/grpc_handler.py b/src/a2a/compat/v0_3/grpc_handler.py index 5f046c765..b7bec26ea 100644 --- a/src/a2a/compat/v0_3/grpc_handler.py +++ b/src/a2a/compat/v0_3/grpc_handler.py @@ -24,7 +24,6 @@ _ERROR_CODE_MAP, DefaultGrpcServerCallContextBuilder, GrpcServerCallContextBuilder, - set_grpc_call_metadata, ) from a2a.server.request_handlers.request_handler import RequestHandler from a2a.utils.errors import A2AError, InvalidParamsError @@ -65,7 +64,6 @@ async def _handle_unary( """Centralized error handling and context management for unary calls.""" try: server_context = self._context_builder.build(context) - set_grpc_call_metadata(server_context, context) result = await handler_func(server_context) except A2AError as e: await self.abort_context(e, context) @@ -81,7 +79,6 @@ async def _handle_stream( """Centralized error handling and context management for streaming calls.""" try: server_context = self._context_builder.build(context) - set_grpc_call_metadata(server_context, context) async for item in handler_func(server_context): yield item except A2AError as e: diff --git a/src/a2a/compat/v0_3/jsonrpc_adapter.py b/src/a2a/compat/v0_3/jsonrpc_adapter.py index a0ed23762..303f5f9f8 100644 --- a/src/a2a/compat/v0_3/jsonrpc_adapter.py +++ b/src/a2a/compat/v0_3/jsonrpc_adapter.py @@ -26,7 +26,7 @@ from a2a.compat.v0_3 import types as types_v03 from a2a.compat.v0_3.context_builders import V03ServerCallContextBuilder from a2a.compat.v0_3.request_handler import RequestHandler03 -from a2a.server.context import TRANSPORT_STATE_KEY, ServerCallContext +from a2a.server.context import ServerCallContext from a2a.server.jsonrpc_models import ( InternalError as CoreInternalError, ) @@ -42,7 +42,6 @@ create_event_source_response, ) from a2a.utils import constants -from a2a.utils.constants import TransportProtocol from a2a.utils.version_validator import validate_version @@ -126,9 +125,6 @@ async def handle_request( ) call_context = self._context_builder.build(request) - call_context.state.setdefault( - TRANSPORT_STATE_KEY, TransportProtocol.JSONRPC.value - ) call_context.tenant = ( getattr(specific_request.params, 'tenant', '') if hasattr(specific_request, 'params') diff --git a/src/a2a/compat/v0_3/rest_adapter.py b/src/a2a/compat/v0_3/rest_adapter.py index 0bd618c44..0c2449dfa 100644 --- a/src/a2a/compat/v0_3/rest_adapter.py +++ b/src/a2a/compat/v0_3/rest_adapter.py @@ -32,14 +32,12 @@ from a2a.compat.v0_3.context_builders import V03ServerCallContextBuilder from a2a.compat.v0_3.rest_handler import REST03Handler -from a2a.server.context import TRANSPORT_STATE_KEY from a2a.server.routes.common import ( DefaultServerCallContextBuilder, ServerCallContextBuilder, create_event_source_response, ) from a2a.utils import json_utils -from a2a.utils.constants import TransportProtocol from a2a.utils.error_handlers import ( rest_error_handler, rest_stream_error_handler, @@ -77,9 +75,6 @@ async def _handle_request( request: Request, ) -> Response: call_context = self._context_builder.build(request) - call_context.state.setdefault( - TRANSPORT_STATE_KEY, TransportProtocol.HTTP_JSON.value - ) response = await method(request, call_context) return JSONResponse(content=response) @@ -97,9 +92,6 @@ async def _handle_streaming_request( ) from e call_context = self._context_builder.build(request) - call_context.state.setdefault( - TRANSPORT_STATE_KEY, TransportProtocol.HTTP_JSON.value - ) async def event_generator( stream: AsyncIterable[Any], diff --git a/src/a2a/plugins/types.py b/src/a2a/plugins/types.py index 22d48a32e..94c5ec373 100644 --- a/src/a2a/plugins/types.py +++ b/src/a2a/plugins/types.py @@ -28,8 +28,8 @@ class ServerRequest: Attributes: method: A2A RPC name, e.g. ``'SendMessage'`` or ``'GetTask'``. payload: The request proto, e.g. ``SendMessageRequest``. - context: The call context. ``state['headers']`` holds the transport - headers and ``state['a2a.transport']`` the transport name. + context: The call context. For HTTP transports ``state['headers']`` + holds the request headers. agent_card: The agent card served by this handler. """ @@ -76,16 +76,12 @@ class ClientRequest: agent_card: The card of the remote agent. context: The call context. ``service_parameters`` holds the outgoing headers. - transport: Protocol binding used for the call, e.g. ``'JSONRPC'``. - url: URL of the selected agent interface. """ method: str payload: Any agent_card: AgentCard context: ClientCallContext - transport: str = '' - url: str = '' @dataclass diff --git a/src/a2a/server/context.py b/src/a2a/server/context.py index c9059014f..833ca44c4 100644 --- a/src/a2a/server/context.py +++ b/src/a2a/server/context.py @@ -10,9 +10,6 @@ State = collections.abc.MutableMapping[str, typing.Any] -TRANSPORT_STATE_KEY = 'a2a.transport' -"""`ServerCallContext.state` key holding the transport of the request.""" - class ServerCallContext(BaseModel): """A context passed when calling a server method. diff --git a/src/a2a/server/request_handlers/grpc_handler.py b/src/a2a/server/request_handlers/grpc_handler.py index bfbcbb6c0..d4d8ed669 100644 --- a/src/a2a/server/request_handlers/grpc_handler.py +++ b/src/a2a/server/request_handlers/grpc_handler.py @@ -27,11 +27,10 @@ HTTP_EXTENSION_HEADER, get_requested_extensions, ) -from a2a.server.context import TRANSPORT_STATE_KEY, ServerCallContext +from a2a.server.context import ServerCallContext from a2a.server.request_handlers.request_handler import RequestHandler from a2a.types import a2a_pb2 from a2a.utils import proto_utils -from a2a.utils.constants import TransportProtocol from a2a.utils.errors import A2A_ERROR_REASONS, A2AError, TaskNotFoundError from a2a.utils.grpc_status import status_to_grpc from a2a.utils.proto_utils import validation_errors_to_bad_request @@ -82,20 +81,6 @@ def _get_metadata_value( ] -def set_grpc_call_metadata( - server_context: ServerCallContext, context: grpc.aio.ServicerContext -) -> None: - """Records the transport and text invocation metadata in the call state.""" - server_context.state.setdefault( - TRANSPORT_STATE_KEY, TransportProtocol.GRPC.value - ) - if 'headers' not in server_context.state: - md = context.invocation_metadata() or () - server_context.state['headers'] = { - k.lower(): v for k, v in md if isinstance(v, str) - } - - _ERROR_CODE_MAP: dict[type[A2AError], grpc.StatusCode] = { types.InvalidRequestError: grpc.StatusCode.INVALID_ARGUMENT, types.MethodNotFoundError: grpc.StatusCode.NOT_FOUND, @@ -439,6 +424,5 @@ def _build_call_context( request: message.Message, ) -> ServerCallContext: server_context = self._context_builder.build(context) - set_grpc_call_metadata(server_context, context) server_context.tenant = getattr(request, 'tenant', '') return server_context diff --git a/src/a2a/server/routes/jsonrpc_dispatcher.py b/src/a2a/server/routes/jsonrpc_dispatcher.py index b4cbaa241..21b71e921 100644 --- a/src/a2a/server/routes/jsonrpc_dispatcher.py +++ b/src/a2a/server/routes/jsonrpc_dispatcher.py @@ -10,7 +10,7 @@ from jsonrpc.jsonrpc2 import JSONRPC20Request, JSONRPC20Response from a2a.compat.v0_3.jsonrpc_adapter import JSONRPC03Adapter -from a2a.server.context import TRANSPORT_STATE_KEY, ServerCallContext +from a2a.server.context import ServerCallContext from a2a.server.events import Event from a2a.server.jsonrpc_models import ( InternalError, @@ -41,7 +41,6 @@ TaskPushNotificationConfig, ) from a2a.utils import constants, json_utils, proto_utils -from a2a.utils.constants import TransportProtocol from a2a.utils.errors import ( A2AError, InvalidParamsError, @@ -309,9 +308,6 @@ async def handle_requests(self, request: Request) -> Response: # noqa: PLR0911, # 3) Build call context and wrap the request for downstream handling call_context = self._context_builder.build(request) - call_context.state.setdefault( - TRANSPORT_STATE_KEY, TransportProtocol.JSONRPC.value - ) call_context.tenant = getattr(specific_request, 'tenant', '') call_context.state['method'] = method call_context.state['request_id'] = request_id diff --git a/src/a2a/server/routes/rest_dispatcher.py b/src/a2a/server/routes/rest_dispatcher.py index 3000b8cf7..54150396e 100644 --- a/src/a2a/server/routes/rest_dispatcher.py +++ b/src/a2a/server/routes/rest_dispatcher.py @@ -5,7 +5,7 @@ from google.protobuf.json_format import MessageToDict, Parse -from a2a.server.context import TRANSPORT_STATE_KEY, ServerCallContext +from a2a.server.context import ServerCallContext from a2a.server.request_handlers.request_handler import RequestHandler from a2a.server.routes.common import ( DefaultServerCallContextBuilder, @@ -20,7 +20,6 @@ SubscribeToTaskRequest, ) from a2a.utils import constants, json_utils, proto_utils -from a2a.utils.constants import TransportProtocol from a2a.utils.error_handlers import ( build_rest_error_payload, rest_error_handler, @@ -107,9 +106,6 @@ def __init__( def _build_call_context(self, request: Request) -> ServerCallContext: call_context = self._context_builder.build(request) - call_context.state.setdefault( - TRANSPORT_STATE_KEY, TransportProtocol.HTTP_JSON.value - ) if 'tenant' in request.path_params: call_context.tenant = request.path_params['tenant'] return call_context diff --git a/tests/client/test_base_client_plugins.py b/tests/client/test_base_client_plugins.py index 5f184995b..066ec701e 100644 --- a/tests/client/test_base_client_plugins.py +++ b/tests/client/test_base_client_plugins.py @@ -69,8 +69,6 @@ def client( transport, interceptors or [], plugins=plugins.sequential(recorder) if recorder else None, - protocol_binding='JSONRPC', - url='http://remote:9000', ) @@ -98,7 +96,7 @@ def restore_global(): @pytest.mark.asyncio -async def test_unary_call_runs_hooks_with_transport_details(): +async def test_unary_call_runs_hooks(): recorder = Recorder() transport = AsyncMock(spec=ClientTransport) transport.get_task.return_value = Task(id='t-1') @@ -110,8 +108,6 @@ async def test_unary_call_runs_hooks_with_transport_details(): assert result.id == 't-1' (request,) = recorder.hooks('before_agent_call') assert request.method == 'GetTask' - assert request.transport == 'JSONRPC' - assert request.url == 'http://remote:9000' assert request.agent_card.name == 'Remote' (response,) = recorder.hooks('after_agent_call') assert response.payload is result @@ -304,7 +300,7 @@ async def test_no_plugins_keeps_context_none(): assert transport.get_task.call_args.kwargs['context'] is None -def test_factory_passes_selected_transport_and_global_plugins(): +def test_factory_clients_use_the_global_manager(): manager = plugins.sequential(Recorder()) plugins.set_global(manager) @@ -314,5 +310,3 @@ def test_factory_passes_selected_transport_and_global_plugins(): assert isinstance(c, BaseClient) assert c._plugins is manager - assert c._protocol_binding == 'JSONRPC' - assert c._url == 'http://remote:9000' diff --git a/tests/integration/test_plugins.py b/tests/integration/test_plugins.py index 86f9cf556..6bab245f2 100644 --- a/tests/integration/test_plugins.py +++ b/tests/integration/test_plugins.py @@ -26,6 +26,7 @@ TaskObserverPlugin, ) from a2a.server.agent_execution import AgentExecutor, RequestContext +from a2a.server.context import ServerCallContext from a2a.server.events import EventQueue from a2a.server.request_handlers import DefaultRequestHandlerV2, GrpcHandler from a2a.server.routes import ( @@ -137,6 +138,13 @@ def send_request(text: str = 'hello') -> SendMessageRequest: ) +def incoming_headers(context: ServerCallContext) -> dict[str, str]: + """Request headers: state['headers'] for HTTP, metadata for gRPC.""" + if 'grpc_context' in context.state: + return dict(context.state['grpc_context'].invocation_metadata()) + return context.state['headers'] + + async def wait_for(predicate: Callable[[], bool], timeout: float = 5) -> None: async def poll() -> None: while not predicate(): @@ -233,17 +241,14 @@ async def test_streaming_message_runs_hooks_on_both_sides(setup: Setup): # Client side. (call,) = setup.caller.hooks('before_agent_call') assert call.method == 'SendStreamingMessage' - assert call.transport == setup.transport - assert call.url client_responses = setup.caller.hooks('after_agent_call') assert len(client_responses) == len(events) + 1 assert client_responses[-1].is_final - # Server side: transport recorded, client header received. + # Server side: the header added by the client plugin arrives. (run,) = setup.server.hooks('before_agent_run') assert run.method == 'SendStreamingMessage' - assert run.context.state['a2a.transport'] == setup.transport - assert run.context.state['headers']['x-trace'] == 'trace-123' + assert incoming_headers(run.context)['x-trace'] == 'trace-123' assert run.agent_card.name == 'Plugin Agent' server_responses = setup.server.hooks('after_agent_run') assert server_responses[-1].is_final @@ -310,7 +315,6 @@ async def test_v03_jsonrpc_requests_run_server_hooks(): assert response.json()['result']['status']['state'] == 'completed' (run,) = recorder.hooks('before_agent_run') assert run.method == 'SendMessage' - assert run.context.state['a2a.transport'] == 'JSONRPC' assert recorder.hooks('after_agent_run')[-1].is_final diff --git a/tests/server/request_handlers/test_grpc_handler.py b/tests/server/request_handlers/test_grpc_handler.py index 8b6ff2713..fa504fc05 100644 --- a/tests/server/request_handlers/test_grpc_handler.py +++ b/tests/server/request_handlers/test_grpc_handler.py @@ -746,28 +746,3 @@ async def mock_stream(*args, **kwargs): server_context = call_args[0][1] assert isinstance(server_context, ServerCallContext) assert server_context.tenant == '' - - -@pytest.mark.asyncio -async def test_call_context_records_transport_and_text_metadata( - grpc_handler: GrpcHandler, - mock_request_handler: AsyncMock, - mock_grpc_context: AsyncMock, -) -> None: - mock_grpc_context.invocation_metadata.return_value = grpc.aio.Metadata( - ('traceparent', '00-abc-def-01'), - ('X-Tenant', 'acme'), - ('trace-bin', b'\x00\x01'), - ) - mock_request_handler.on_get_task.return_value = types.Task(id='t-1') - - await grpc_handler.GetTask( - a2a_pb2.GetTaskRequest(id='t-1'), mock_grpc_context - ) - - call_context = mock_request_handler.on_get_task.call_args[0][1] - assert call_context.state['a2a.transport'] == 'GRPC' - assert call_context.state['headers'] == { - 'traceparent': '00-abc-def-01', - 'x-tenant': 'acme', - } From f50a7fcabe487be0af86117643e150661a250e4f Mon Sep 17 00:00:00 2001 From: Mykyta Netipa Date: Sat, 17 Oct 2026 18:59:31 +0000 Subject: [PATCH 07/11] refactor(client): pass plugin RPC names at each call site Replace the lookup table that mapped client method names to A2A RPC names with an rpc_method argument passed by each BaseClient method, the same way DefaultRequestHandlerV2 names its methods in @server_plugin_hooks. Interceptors keep receiving the existing method strings; rpc_method is keyword-only and defaults to method, so callers of the private helpers are unaffected. --- src/a2a/client/base_client.py | 56 +++++++++------ tests/client/test_base_client_plugins.py | 91 ++++++++++++++++++++++++ 2 files changed, 127 insertions(+), 20 deletions(-) diff --git a/src/a2a/client/base_client.py b/src/a2a/client/base_client.py index 4637416b9..6c4afd31d 100644 --- a/src/a2a/client/base_client.py +++ b/src/a2a/client/base_client.py @@ -34,22 +34,6 @@ ) -# Client method names mapped to the A2A RPC names reported to plugins. -_RPC_NAMES = { - 'send_message': 'SendMessage', - 'send_message_streaming': 'SendStreamingMessage', - 'get_task': 'GetTask', - 'list_tasks': 'ListTasks', - 'cancel_task': 'CancelTask', - 'create_task_push_notification_config': 'CreateTaskPushNotificationConfig', - 'get_task_push_notification_config': 'GetTaskPushNotificationConfig', - 'list_task_push_notification_configs': 'ListTaskPushNotificationConfigs', - 'delete_task_push_notification_config': 'DeleteTaskPushNotificationConfig', - 'subscribe': 'SubscribeToTask', - 'get_extended_agent_card': 'GetExtendedAgentCard', -} - - class BaseClient(Client): """Base implementation of the A2A client, containing transport-independent logic. @@ -107,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 @@ -132,6 +117,7 @@ async def send_message( 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) @@ -193,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 @@ -209,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 @@ -233,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 @@ -257,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( @@ -283,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( @@ -309,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( @@ -332,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( @@ -369,6 +362,7 @@ async def subscribe( 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 @@ -401,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) @@ -417,10 +412,13 @@ async def close(self) -> None: await self._transport.close() def _client_request( - self, input_data: Any, method: str, context: ClientCallContext | None + self, + input_data: Any, + rpc_method: str, + context: ClientCallContext | None, ) -> ClientRequest: return ClientRequest( - method=_RPC_NAMES.get(method, method), + method=rpc_method, payload=input_data, agent_card=self._card, context=context if context is not None else ClientCallContext(), @@ -434,14 +432,23 @@ 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, method, context) + request = self._client_request( + input_data, rpc_method or method, context + ) await plugins.before_agent_call(request) try: result = await self._call_through_interceptors( @@ -507,7 +514,14 @@ async def _execute_stream_with_interceptors( transport_call: Callable[ [Any, ClientCallContext | None], 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( @@ -519,7 +533,9 @@ async def _execute_stream_with_interceptors( yield event return - request = self._client_request(input_data, method, context) + request = self._client_request( + input_data, rpc_method or method, context + ) await plugins.before_agent_call(request) error: BaseException | None = None try: diff --git a/tests/client/test_base_client_plugins.py b/tests/client/test_base_client_plugins.py index 066ec701e..dd80d3959 100644 --- a/tests/client/test_base_client_plugins.py +++ b/tests/client/test_base_client_plugins.py @@ -15,14 +15,22 @@ AgentCapabilities, AgentCard, AgentInterface, + CancelTaskRequest, + DeleteTaskPushNotificationConfigRequest, + GetExtendedAgentCardRequest, + GetTaskPushNotificationConfigRequest, GetTaskRequest, + ListTaskPushNotificationConfigsRequest, + ListTasksRequest, Message, Part, Role, SendMessageRequest, SendMessageResponse, StreamResponse, + SubscribeToTaskRequest, Task, + TaskPushNotificationConfig, TaskState, TaskStatus, TaskStatusUpdateEvent, @@ -310,3 +318,86 @@ def test_factory_clients_use_the_global_manager(): assert isinstance(c, BaseClient) assert c._plugins is manager + + +@pytest.mark.parametrize( + ('call', 'rpc_name'), + [ + (lambda c: c.list_tasks(ListTasksRequest()), 'ListTasks'), + (lambda c: c.cancel_task(CancelTaskRequest(id='t-1')), 'CancelTask'), + ( + lambda c: c.create_task_push_notification_config( + TaskPushNotificationConfig(task_id='t-1', url='http://hook') + ), + 'CreateTaskPushNotificationConfig', + ), + ( + lambda c: c.get_task_push_notification_config( + GetTaskPushNotificationConfigRequest(task_id='t-1', id='c-1') + ), + 'GetTaskPushNotificationConfig', + ), + ( + lambda c: c.list_task_push_notification_configs( + ListTaskPushNotificationConfigsRequest(task_id='t-1') + ), + 'ListTaskPushNotificationConfigs', + ), + ( + lambda c: c.delete_task_push_notification_config( + DeleteTaskPushNotificationConfigRequest(task_id='t-1', id='c-1') + ), + 'DeleteTaskPushNotificationConfig', + ), + ( + lambda c: c.get_extended_agent_card(GetExtendedAgentCardRequest()), + 'GetExtendedAgentCard', + ), + ], +) +@pytest.mark.asyncio +async def test_unary_methods_report_rpc_names(call, rpc_name): + recorder = Recorder() + + await call(client(recorder, AsyncMock(spec=ClientTransport))) + + assert [r.method for r in recorder.hooks('before_agent_call')] == [rpc_name] + assert recorder.hooks('after_agent_call')[-1].method == rpc_name + + +@pytest.mark.asyncio +async def test_subscribe_reports_rpc_name(): + recorder = Recorder() + transport = AsyncMock(spec=ClientTransport) + + async def stream(request, *, context=None): + yield status_event(TaskState.TASK_STATE_WORKING) + + transport.subscribe = stream + + async for _ in client(recorder, transport).subscribe( + SubscribeToTaskRequest(id='t-1') + ): + pass + + assert recorder.hooks('before_agent_call')[0].method == 'SubscribeToTask' + assert all( + r.method == 'SubscribeToTask' + for r in recorder.hooks('after_agent_call') + ) + + +@pytest.mark.asyncio +async def test_rpc_method_defaults_to_method(): + recorder = Recorder() + transport = AsyncMock(spec=ClientTransport) + c = client(recorder, transport) + + await c._execute_with_interceptors( + input_data=GetTaskRequest(id='t-1'), + method='custom_call', + context=None, + transport_call=AsyncMock(return_value=Task(id='t-1')), + ) + + assert recorder.hooks('before_agent_call')[0].method == 'custom_call' From d8ab490d0275e894e58c7df28c49301ce6ddcdd6 Mon Sep 17 00:00:00 2001 From: Mykyta Netipa Date: Sat, 17 Oct 2026 19:12:37 +0000 Subject: [PATCH 08/11] refactor(server): let the push sender own its plugin manager BasePushNotificationSender now takes a plugins argument (global manager by default) and runs push hooks for every delivery, whichever request handler uses it. This replaces the contextvar that let only DefaultRequestHandlerV2 enable push hooks: not wiring the legacy handler for plugins does not require shared components to block them. ActiveTask is back to its state on main, and a2a.plugins._runtime is removed. --- src/a2a/plugins/_runtime.py | 40 ------------ src/a2a/plugins/manager.py | 7 ++- src/a2a/server/agent_execution/active_task.py | 13 +--- .../agent_execution/active_task_registry.py | 1 - .../tasks/base_push_notification_sender.py | 11 +++- ...test_default_request_handler_v2_plugins.py | 61 ++++++++++++++++--- 6 files changed, 68 insertions(+), 65 deletions(-) delete mode 100644 src/a2a/plugins/_runtime.py diff --git a/src/a2a/plugins/_runtime.py b/src/a2a/plugins/_runtime.py deleted file mode 100644 index b3fe4378d..000000000 --- a/src/a2a/plugins/_runtime.py +++ /dev/null @@ -1,40 +0,0 @@ -"""Internal plumbing for scoping plugin hooks to `DefaultRequestHandlerV2`.""" - -from __future__ import annotations - -import contextlib - -from contextvars import ContextVar -from typing import TYPE_CHECKING - - -if TYPE_CHECKING: - from collections.abc import Iterator - - from a2a.plugins.manager import PluginManager - - -# Set by `ActiveTask` around push delivery, so `BasePushNotificationSender` -# runs push hooks for the V2 path only. Child tasks created inside the scope -# (e.g. by asyncio.gather) inherit the value. -_push_plugins: ContextVar[PluginManager | None] = ContextVar( - 'a2a_push_plugins', default=None -) - - -@contextlib.contextmanager -def push_plugins_scope(manager: PluginManager | None) -> Iterator[None]: - """Makes ``manager`` the push plugin manager inside the block.""" - token = _push_plugins.set(manager) - try: - yield - finally: - _push_plugins.reset(token) - - -def current_push_plugins() -> PluginManager | None: - """Returns the manager set by `push_plugins_scope`, if it has push hooks.""" - manager = _push_plugins.get() - if manager is None or not manager.has_push_hooks: - return None - return manager diff --git a/src/a2a/plugins/manager.py b/src/a2a/plugins/manager.py index c115641da..145b6dcd7 100644 --- a/src/a2a/plugins/manager.py +++ b/src/a2a/plugins/manager.py @@ -272,10 +272,11 @@ def sequential(*plugins: Plugin) -> SequentialPluginManager: def set_global(manager: PluginManager) -> None: - """Sets the manager used by handlers and clients created afterwards. + """Sets the manager used by handlers, clients and push senders created afterwards. - Handlers and clients capture the manager when they are constructed, so - call this once at startup, before creating them. + `DefaultRequestHandlerV2`, clients and `BasePushNotificationSender` + capture the manager when they are constructed, so call this once at + startup, before creating them. """ global _global_manager # noqa: PLW0603 if not isinstance(manager, PluginManager): diff --git a/src/a2a/server/agent_execution/active_task.py b/src/a2a/server/agent_execution/active_task.py index bb6530b91..5e3dab50f 100644 --- a/src/a2a/server/agent_execution/active_task.py +++ b/src/a2a/server/agent_execution/active_task.py @@ -48,7 +48,6 @@ if TYPE_CHECKING: from collections.abc import AsyncGenerator, Callable - from a2a.plugins.manager import PluginManager from a2a.server.agent_execution.agent_executor import AgentExecutor from a2a.server.cluster.event_stream import TaskEventStream from a2a.server.context import ServerCallContext @@ -57,7 +56,6 @@ ) from a2a.server.tasks.task_manager import TaskManager -from a2a.plugins._runtime import push_plugins_scope from a2a.server.cluster.event_stream import VersionedEvent from a2a.server.cluster.task_store import ConcurrentTaskModificationError from a2a.server.events.event_queue_v2 import ( @@ -361,10 +359,9 @@ async def _update_task_state( self.active_task._task_id, ) try: - with push_plugins_scope(self.active_task._plugins): - await self.active_task._push_sender.send_notification( - self.active_task._task_id, event - ) + await self.active_task._push_sender.send_notification( + self.active_task._task_id, event + ) except Exception: # Push delivery is best-effort: a sender failure must not # fail the task or disturb the event stream. This guard is @@ -444,7 +441,6 @@ def __init__( # noqa: PLR0913 push_sender: PushNotificationSender | None = None, on_cleanup: Callable[[ActiveTask], None] | None = None, event_stream: TaskEventStream | None = None, - plugins: PluginManager | None = None, ) -> None: """Initializes the ActiveTask. @@ -458,8 +454,6 @@ def __init__( # noqa: PLR0913 the task from the ActiveTaskRegistry. event_stream: Optional cross-replica stream; applied events are published to it so other replicas can observe them. - plugins: Optional plugin manager whose push hooks run when - `push_sender` delivers notifications. """ # --- Core Dependencies --- self._agent_executor = agent_executor @@ -472,7 +466,6 @@ def __init__( # noqa: PLR0913 self._push_sender = push_sender self._on_cleanup = on_cleanup self._event_stream = event_stream - self._plugins = plugins # --- Synchronization Primitives --- # `_lock` protects structural lifecycle changes: start(), subscribe() counting, diff --git a/src/a2a/server/agent_execution/active_task_registry.py b/src/a2a/server/agent_execution/active_task_registry.py index 11441aa35..b2f7b88c3 100644 --- a/src/a2a/server/agent_execution/active_task_registry.py +++ b/src/a2a/server/agent_execution/active_task_registry.py @@ -77,7 +77,6 @@ async def get_or_create( push_sender=self._push_sender, on_cleanup=self._on_active_task_cleanup, event_stream=self._event_stream, - plugins=self._plugins, ) self._active_tasks[task_id] = active_task diff --git a/src/a2a/server/tasks/base_push_notification_sender.py b/src/a2a/server/tasks/base_push_notification_sender.py index b425251c4..bade95752 100644 --- a/src/a2a/server/tasks/base_push_notification_sender.py +++ b/src/a2a/server/tasks/base_push_notification_sender.py @@ -8,8 +8,7 @@ from google.protobuf.json_format import MessageToDict -from a2a.plugins._runtime import current_push_plugins -from a2a.plugins.manager import PluginManager +from a2a.plugins.manager import PluginManager, get_global from a2a.plugins.types import PushRequest, PushResult from a2a.server.context import ServerCallContext from a2a.server.tasks.push_notification_config_store import ( @@ -36,6 +35,7 @@ def __init__( context: ServerCallContext | None = None, *, push_url_validator: Callable[[str], Awaitable[bool]] | None = None, + plugins: PluginManager | None = None, ) -> None: """Initializes the BasePushNotificationSender. @@ -54,6 +54,10 @@ def __init__( None (no library screening). The spec lists these checks as SHOULD, so deployments that want the built-in policy should pass ``validate_push_notification_url``. + plugins: Plugin manager whose `PushPlugin` hooks run for every + delivery, whichever request handler uses this sender. + Defaults to the global manager at construction time; it is + independent of the manager passed to the request handler. """ if context is not None: logger.warning( @@ -68,6 +72,7 @@ def __init__( self._client = httpx_client self._config_store = config_store self._push_url_validator = push_url_validator + self._plugins = plugins if plugins is not None else get_global() async def send_notification( self, task_id: str, event: PushNotificationEvent @@ -125,7 +130,7 @@ async def _dispatch_notification( return False if not accepted: return False - plugins = current_push_plugins() + plugins = self._plugins if self._plugins.has_push_hooks else None request = PushRequest( task_id, url, self._build_headers(push_info, task_id), event ) diff --git a/tests/server/request_handlers/test_default_request_handler_v2_plugins.py b/tests/server/request_handlers/test_default_request_handler_v2_plugins.py index 1e59d3972..476d55375 100644 --- a/tests/server/request_handlers/test_default_request_handler_v2_plugins.py +++ b/tests/server/request_handlers/test_default_request_handler_v2_plugins.py @@ -510,6 +510,7 @@ async def test_remote_cancel_reports_task_event(): def push_setup( responses: dict[str, int], + manager: plugins.PluginManager | None = None, ) -> tuple[ InMemoryPushNotificationConfigStore, BasePushNotificationSender, @@ -523,7 +524,9 @@ def handle(request: httpx.Request) -> httpx.Response: config_store = InMemoryPushNotificationConfigStore() sender = BasePushNotificationSender( - httpx.AsyncClient(transport=httpx.MockTransport(handle)), config_store + httpx.AsyncClient(transport=httpx.MockTransport(handle)), + config_store, + plugins=manager, ) return config_store, sender, received @@ -540,8 +543,10 @@ def push_request(url: str) -> SendMessageRequest: @pytest.mark.asyncio async def test_push_hooks_run_for_v2_deliveries(): - config_store, sender, received = push_setup({'http://hook/fail': 500}) recorder = Recorder() + config_store, sender, received = push_setup( + {'http://hook/fail': 500}, plugins.sequential(recorder) + ) handler = make_handler( recorder, push_config_store=config_store, push_sender=sender ) @@ -590,14 +595,15 @@ async def on_push_sent(self, result): (Fail(), lambda r: isinstance(r.error, RuntimeError)), ): seen.clear() - config_store, sender, received = push_setup({}) + config_store, sender, received = push_setup( + {}, plugins.sequential(plugin, Seen()) + ) handler = DefaultRequestHandlerV2( agent_executor=TaskAgent(), task_store=InMemoryTaskStore(), agent_card=CARD, push_config_store=config_store, push_sender=sender, - plugins=plugins.sequential(plugin, Seen()), ) await handler.on_message_send( push_request('http://hook/x'), ServerCallContext() @@ -613,6 +619,7 @@ async def reject(url: str) -> bool: return False received: list[httpx.Request] = [] + recorder = Recorder() config_store = InMemoryPushNotificationConfigStore() sender = BasePushNotificationSender( httpx.AsyncClient( @@ -622,8 +629,8 @@ async def reject(url: str) -> bool: ), config_store, push_url_validator=reject, + plugins=plugins.sequential(recorder), ) - recorder = Recorder() handler = make_handler( recorder, push_config_store=config_store, push_sender=sender ) @@ -639,7 +646,8 @@ async def reject(url: str) -> bool: @pytest.mark.asyncio -async def test_legacy_handler_runs_no_plugin_hooks(): +async def test_legacy_handler_runs_only_push_hooks(): + """The legacy handler is not wired for plugins; the sender is on its own.""" recorder = Recorder() plugins.set_global(plugins.sequential(recorder)) config_store, sender, received = push_setup({}) @@ -654,9 +662,46 @@ async def test_legacy_handler_runs_no_plugin_hooks(): result = await handler.on_message_send( push_request('http://hook/legacy'), ServerCallContext() ) - await asyncio.sleep(0.1) + await wait_for(lambda: bool(recorder.hooks('on_push_sent'))) assert result.status.state == TaskState.TASK_STATE_COMPLETED + assert received + assert all(r.headers['x-plugin'] == 'yes' for r in received) + assert set(recorder.names()) == {'before_push_send', 'on_push_sent'} + + +def test_push_sender_uses_global_manager_unless_given_one(): + manager = plugins.sequential(Recorder()) + plugins.set_global(manager) + other = plugins.sequential() + + default = BasePushNotificationSender( + httpx.AsyncClient(), InMemoryPushNotificationConfigStore() + ) + explicit = BasePushNotificationSender( + httpx.AsyncClient(), + InMemoryPushNotificationConfigStore(), + plugins=other, + ) + + assert default._plugins is manager + assert explicit._plugins is other + + +@pytest.mark.asyncio +async def test_handler_plugins_do_not_reach_a_sender_with_its_own_manager(): + recorder = Recorder() + config_store, sender, received = push_setup({}) # global: no plugins + handler = make_handler( + recorder, push_config_store=config_store, push_sender=sender + ) + + await handler.on_message_send( + push_request('http://hook/x'), ServerCallContext() + ) + await wait_for(lambda: 'after_execution' in recorder.names()) + await asyncio.sleep(0.05) + assert received assert 'x-plugin' not in received[0].headers - assert recorder.calls == [] + assert recorder.hooks('before_push_send') == [] From a6ed492233e1ee0862ab4c985b78b8e199fcbe8b Mon Sep 17 00:00:00 2001 From: Mykyta Netipa Date: Thu, 22 Oct 2026 10:21:20 +0000 Subject: [PATCH 09/11] refactor(plugins): rename server and client hooks ServerPlugin's before/after_agent_run run for every request handler method, including ones that never run an agent (task retrieval, push config CRUD), so they become before/after_handler_run. ClientPlugin's before/after_agent_call likewise run for every client request and become before/after_client_request, matching ClientRequest and ClientResponse. --- src/a2a/client/base_client.py | 16 ++-- src/a2a/plugins/base.py | 12 +-- src/a2a/plugins/manager.py | 26 +++--- src/a2a/plugins/types.py | 12 +-- .../server/request_handlers/_plugin_hooks.py | 16 ++-- tests/client/test_base_client_plugins.py | 52 +++++------ tests/integration/test_plugins.py | 36 ++++---- tests/plugins/test_manager.py | 86 +++++++++---------- ...test_default_request_handler_v2_plugins.py | 30 +++---- 9 files changed, 147 insertions(+), 139 deletions(-) diff --git a/src/a2a/client/base_client.py b/src/a2a/client/base_client.py index 6c4afd31d..b13873a52 100644 --- a/src/a2a/client/base_client.py +++ b/src/a2a/client/base_client.py @@ -37,8 +37,8 @@ class BaseClient(Client): """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. + Client plugin hooks run around every call: ``before_client_request`` before + the interceptors, ``after_client_request`` after them. """ def __init__( @@ -449,17 +449,17 @@ async def _execute_with_interceptors( request = self._client_request( input_data, rpc_method or method, context ) - await plugins.before_agent_call(request) + 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_agent_call( + await plugins.after_client_request( ClientResponse(request.method, None, request.context, error=e) ) raise - await plugins.after_agent_call( + await plugins.after_client_request( ClientResponse(request.method, result, request.context) ) return result @@ -536,7 +536,7 @@ async def _execute_stream_with_interceptors( request = self._client_request( input_data, rpc_method or method, context ) - await plugins.before_agent_call(request) + await plugins.before_client_request(request) error: BaseException | None = None try: async with aclosing( @@ -545,7 +545,7 @@ async def _execute_stream_with_interceptors( ) ) as events: async for event in events: - await plugins.after_agent_call( + await plugins.after_client_request( ClientResponse( request.method, event, @@ -560,7 +560,7 @@ async def _execute_stream_with_interceptors( raise finally: # Also runs when the caller stops iterating early. - await plugins.after_agent_call( + await plugins.after_client_request( ClientResponse( request.method, None, request.context, error=error ) diff --git a/src/a2a/plugins/base.py b/src/a2a/plugins/base.py index 4e3e391ae..23a51eb53 100644 --- a/src/a2a/plugins/base.py +++ b/src/a2a/plugins/base.py @@ -37,23 +37,23 @@ class Plugin: class ServerPlugin(Plugin): """Hooks around each request handled by `DefaultRequestHandlerV2`.""" - async def before_agent_run(self, request: ServerRequest) -> None: + async def before_handler_run(self, request: ServerRequest) -> None: """Called before the handler processes a request. Raise to reject the request; the exception is returned to the caller. """ - async def after_agent_run(self, response: ServerResponse) -> None: + async def after_handler_run(self, response: ServerResponse) -> None: """Called with the handler result, per stream event, and at the end.""" class ClientPlugin(Plugin): - """Hooks around each call made by an A2A `Client`.""" + """Hooks around each request made by an A2A `Client`.""" - async def before_agent_call(self, request: ClientRequest) -> None: - """Called before the call is sent. Raise to cancel it.""" + async def before_client_request(self, request: ClientRequest) -> None: + """Called before the request is sent. Raise to cancel it.""" - async def after_agent_call(self, response: ClientResponse) -> None: + async def after_client_request(self, response: ClientResponse) -> None: """Called with the result, per stream event, and at the end.""" diff --git a/src/a2a/plugins/manager.py b/src/a2a/plugins/manager.py index 145b6dcd7..0b2af80d9 100644 --- a/src/a2a/plugins/manager.py +++ b/src/a2a/plugins/manager.py @@ -157,35 +157,37 @@ def has_push_hooks(self) -> bool: # noqa: D102 # --- server --- - async def before_agent_run(self, request: ServerRequest) -> None: # noqa: D102 + async def before_handler_run(self, request: ServerRequest) -> None: # noqa: D102 ran: list[ServerPlugin] = [] for plugin in self._server: try: - await plugin.before_agent_run(request) + await plugin.before_handler_run(request) except Exception as e: response = ServerResponse( request.method, None, request.context, error=e ) for done in reversed(ran): await _run_safely( - done, 'after_agent_run', done.after_agent_run(response) + done, + 'after_handler_run', + done.after_handler_run(response), ) raise ran.append(plugin) - async def after_agent_run(self, response: ServerResponse) -> None: # noqa: D102 + async def after_handler_run(self, response: ServerResponse) -> None: # noqa: D102 for plugin in reversed(self._server): await _run_safely( - plugin, 'after_agent_run', plugin.after_agent_run(response) + plugin, 'after_handler_run', plugin.after_handler_run(response) ) # --- client --- - async def before_agent_call(self, request: ClientRequest) -> None: # noqa: D102 + async def before_client_request(self, request: ClientRequest) -> None: # noqa: D102 ran: list[ClientPlugin] = [] for plugin in self._client: try: - await plugin.before_agent_call(request) + await plugin.before_client_request(request) except Exception as e: response = ClientResponse( request.method, None, request.context, error=e @@ -193,16 +195,18 @@ async def before_agent_call(self, request: ClientRequest) -> None: # noqa: D102 for done in reversed(ran): await _run_safely( done, - 'after_agent_call', - done.after_agent_call(response), + 'after_client_request', + done.after_client_request(response), ) raise ran.append(plugin) - async def after_agent_call(self, response: ClientResponse) -> None: # noqa: D102 + async def after_client_request(self, response: ClientResponse) -> None: # noqa: D102 for plugin in reversed(self._client): await _run_safely( - plugin, 'after_agent_call', plugin.after_agent_call(response) + plugin, + 'after_client_request', + plugin.after_client_request(response), ) # --- execution --- diff --git a/src/a2a/plugins/types.py b/src/a2a/plugins/types.py index 94c5ec373..5fbb9f746 100644 --- a/src/a2a/plugins/types.py +++ b/src/a2a/plugins/types.py @@ -23,7 +23,7 @@ @dataclass class ServerRequest: - """An incoming request, passed to `ServerPlugin.before_agent_run`. + """An incoming request, passed to `ServerPlugin.before_handler_run`. Attributes: method: A2A RPC name, e.g. ``'SendMessage'`` or ``'GetTask'``. @@ -41,7 +41,7 @@ class ServerRequest: @dataclass class ServerResponse: - """A handler result, passed to `ServerPlugin.after_agent_run`. + """A handler result, passed to `ServerPlugin.after_handler_run`. For a unary call the hook runs once with ``is_final=True``. For a stream it runs once per event with ``is_stream_event=True, is_final=False`` and @@ -52,7 +52,7 @@ class ServerResponse: method: A2A RPC name. payload: The handler result or stream event, ``None`` on error or stream end. - context: The same call context passed to `before_agent_run`. + context: The same call context passed to `before_handler_run`. error: The exception raised by the handler, if any. is_stream_event: Whether ``payload`` is one event of a stream. is_final: Whether this is the last call for the request. @@ -68,7 +68,7 @@ class ServerResponse: @dataclass class ClientRequest: - """An outgoing call, passed to `ClientPlugin.before_agent_call`. + """An outgoing call, passed to `ClientPlugin.before_client_request`. Attributes: method: A2A RPC name, matching the server side. @@ -86,7 +86,7 @@ class ClientRequest: @dataclass class ClientResponse: - """A call result, passed to `ClientPlugin.after_agent_call`. + """A call result, passed to `ClientPlugin.after_client_request`. Follows the same ``is_stream_event`` / ``is_final`` rules as `ServerResponse`. @@ -94,7 +94,7 @@ class ClientResponse: Attributes: method: A2A RPC name. payload: The response or stream event, ``None`` on error or stream end. - context: The same call context passed to `before_agent_call`. + context: The same call context passed to `before_client_request`. error: The exception raised by the call, if any. is_stream_event: Whether ``payload`` is one event of a stream. is_final: Whether this is the last call for the request. diff --git a/src/a2a/server/request_handlers/_plugin_hooks.py b/src/a2a/server/request_handlers/_plugin_hooks.py index 5774e08b9..adf0c02a3 100644 --- a/src/a2a/server/request_handlers/_plugin_hooks.py +++ b/src/a2a/server/request_handlers/_plugin_hooks.py @@ -33,10 +33,10 @@ def agent_card(self) -> AgentCard: ... def server_plugin_hooks(method: str) -> Callable[[_F], _F]: - """Runs `before_agent_run` / `after_agent_run` around a handler method. + """Runs `before_handler_run` / `after_handler_run` around a handler method. Apply as the outermost decorator, so that validation errors raised by - inner decorators also reach `after_agent_run`. When the handler's manager + inner decorators also reach `after_handler_run`. When the handler's manager has no server hooks the method runs unchanged. Args: @@ -63,7 +63,7 @@ async def stream_wrapper( yield event return - await plugins.before_agent_run( + await plugins.before_handler_run( ServerRequest(method, params, context, self.agent_card) ) error: BaseException | None = None @@ -72,7 +72,7 @@ async def stream_wrapper( func(self, params, context, *args, **kwargs) ) as events: async for event in events: - await plugins.after_agent_run( + await plugins.after_handler_run( ServerResponse( method, event, @@ -87,7 +87,7 @@ async def stream_wrapper( raise finally: # Also runs when the client disconnects (aclose()). - await plugins.after_agent_run( + await plugins.after_handler_run( ServerResponse(method, None, context, error=error) ) @@ -105,17 +105,17 @@ async def unary_wrapper( if not plugins.has_server_hooks: return await func(self, params, context, *args, **kwargs) - await plugins.before_agent_run( + await plugins.before_handler_run( ServerRequest(method, params, context, self.agent_card) ) try: result = await func(self, params, context, *args, **kwargs) except Exception as e: - await plugins.after_agent_run( + await plugins.after_handler_run( ServerResponse(method, None, context, error=e) ) raise - await plugins.after_agent_run( + await plugins.after_handler_run( ServerResponse(method, result, context) ) return result diff --git a/tests/client/test_base_client_plugins.py b/tests/client/test_base_client_plugins.py index dd80d3959..8a1ed4e09 100644 --- a/tests/client/test_base_client_plugins.py +++ b/tests/client/test_base_client_plugins.py @@ -44,14 +44,14 @@ def __init__(self) -> None: def hooks(self, name: str) -> list[Any]: return [arg for hook, arg in self.calls if hook == name] - async def before_agent_call(self, request: ClientRequest) -> None: + async def before_client_request(self, request: ClientRequest) -> None: params = dict(request.context.service_parameters or {}) params['x-trace'] = 'abc' request.context.service_parameters = params - self.calls.append(('before_agent_call', request)) + self.calls.append(('before_client_request', request)) - async def after_agent_call(self, response: ClientResponse) -> None: - self.calls.append(('after_agent_call', response)) + async def after_client_request(self, response: ClientResponse) -> None: + self.calls.append(('after_client_request', response)) def card(streaming: bool = True) -> AgentCard: @@ -114,10 +114,10 @@ async def test_unary_call_runs_hooks(): ) assert result.id == 't-1' - (request,) = recorder.hooks('before_agent_call') + (request,) = recorder.hooks('before_client_request') assert request.method == 'GetTask' assert request.agent_card.name == 'Remote' - (response,) = recorder.hooks('after_agent_call') + (response,) = recorder.hooks('after_client_request') assert response.payload is result assert response.is_final assert response.context is request.context @@ -142,8 +142,8 @@ async def test_non_streaming_send_message_is_a_unary_call(): ] assert len(events) == 1 - assert recorder.hooks('before_agent_call')[0].method == 'SendMessage' - (response,) = recorder.hooks('after_agent_call') + assert recorder.hooks('before_client_request')[0].method == 'SendMessage' + (response,) = recorder.hooks('after_client_request') assert response.is_final @@ -168,9 +168,9 @@ async def stream(request, *, context=None): ] assert events == streamed - (request,) = recorder.hooks('before_agent_call') + (request,) = recorder.hooks('before_client_request') assert request.method == 'SendStreamingMessage' - responses = recorder.hooks('after_agent_call') + responses = recorder.hooks('after_client_request') assert [r.payload for r in responses[:-1]] == streamed assert all(r.is_stream_event and not r.is_final for r in responses[:-1]) assert responses[-1].is_final @@ -192,11 +192,11 @@ async def stream(request, *, context=None): await anext(events) await events.aclose() - assert recorder.hooks('after_agent_call')[-1].is_final + assert recorder.hooks('after_client_request')[-1].is_final @pytest.mark.asyncio -async def test_transport_errors_reach_after_agent_call(): +async def test_transport_errors_reach_after_client_request(): recorder = Recorder() transport = AsyncMock(spec=ClientTransport) transport.get_task.side_effect = ConnectionError('down') @@ -204,7 +204,7 @@ async def test_transport_errors_reach_after_agent_call(): with pytest.raises(ConnectionError): await client(recorder, transport).get_task(GetTaskRequest(id='t-1')) - (response,) = recorder.hooks('after_agent_call') + (response,) = recorder.hooks('after_client_request') assert isinstance(response.error, ConnectionError) async def failing(request, *, context=None): @@ -217,7 +217,7 @@ async def failing(request, *, context=None): async for _ in client(recorder, transport).send_message(send_request()): pass assert isinstance( - recorder.hooks('after_agent_call')[-1].error, ConnectionError + recorder.hooks('after_client_request')[-1].error, ConnectionError ) @@ -226,10 +226,10 @@ async def test_plugins_wrap_interceptors_including_early_return(): order: list[str] = [] class Order(plugins.ClientPlugin): - async def before_agent_call(self, request): + async def before_client_request(self, request): order.append('plugin.before') - async def after_agent_call(self, response): + async def after_client_request(self, response): order.append('plugin.after') class Cached(ClientCallInterceptor): @@ -262,9 +262,9 @@ async def after(self, args: AfterArgs) -> None: @pytest.mark.asyncio -async def test_before_agent_call_can_cancel_the_call(): +async def test_before_client_request_can_cancel_the_call(): class Deny(plugins.ClientPlugin): - async def before_agent_call(self, request): + async def before_client_request(self, request): raise PermissionError('no') transport = AsyncMock(spec=ClientTransport) @@ -295,7 +295,7 @@ async def test_existing_context_is_passed_through(): GetTaskRequest(id='t-1'), context=context ) - assert recorder.hooks('before_agent_call')[0].context is context + assert recorder.hooks('before_client_request')[0].context is context @pytest.mark.asyncio @@ -361,8 +361,10 @@ async def test_unary_methods_report_rpc_names(call, rpc_name): await call(client(recorder, AsyncMock(spec=ClientTransport))) - assert [r.method for r in recorder.hooks('before_agent_call')] == [rpc_name] - assert recorder.hooks('after_agent_call')[-1].method == rpc_name + assert [r.method for r in recorder.hooks('before_client_request')] == [ + rpc_name + ] + assert recorder.hooks('after_client_request')[-1].method == rpc_name @pytest.mark.asyncio @@ -380,10 +382,12 @@ async def stream(request, *, context=None): ): pass - assert recorder.hooks('before_agent_call')[0].method == 'SubscribeToTask' + assert ( + recorder.hooks('before_client_request')[0].method == 'SubscribeToTask' + ) assert all( r.method == 'SubscribeToTask' - for r in recorder.hooks('after_agent_call') + for r in recorder.hooks('after_client_request') ) @@ -400,4 +404,4 @@ async def test_rpc_method_defaults_to_method(): transport_call=AsyncMock(return_value=Task(id='t-1')), ) - assert recorder.hooks('before_agent_call')[0].method == 'custom_call' + assert recorder.hooks('before_client_request')[0].method == 'custom_call' diff --git a/tests/integration/test_plugins.py b/tests/integration/test_plugins.py index 6bab245f2..16a9d1056 100644 --- a/tests/integration/test_plugins.py +++ b/tests/integration/test_plugins.py @@ -62,20 +62,20 @@ def __init__(self) -> None: def hooks(self, name: str) -> list[Any]: return [arg for hook, arg in self.calls if hook == name] - async def before_agent_run(self, request: ServerRequest) -> None: - self.calls.append(('before_agent_run', request)) + async def before_handler_run(self, request: ServerRequest) -> None: + self.calls.append(('before_handler_run', request)) - async def after_agent_run(self, response: ServerResponse) -> None: - self.calls.append(('after_agent_run', response)) + async def after_handler_run(self, response: ServerResponse) -> None: + self.calls.append(('after_handler_run', response)) - async def before_agent_call(self, request: ClientRequest) -> None: + async def before_client_request(self, request: ClientRequest) -> None: params = dict(request.context.service_parameters or {}) params['x-trace'] = 'trace-123' request.context.service_parameters = params - self.calls.append(('before_agent_call', request)) + self.calls.append(('before_client_request', request)) - async def after_agent_call(self, response: ClientResponse) -> None: - self.calls.append(('after_agent_call', response)) + async def after_client_request(self, response: ClientResponse) -> None: + self.calls.append(('after_client_request', response)) async def before_execution(self, request_context: RequestContext) -> None: self.calls.append(('before_execution', request_context)) @@ -239,18 +239,18 @@ async def test_streaming_message_runs_hooks_on_both_sides(setup: Setup): ) # Client side. - (call,) = setup.caller.hooks('before_agent_call') + (call,) = setup.caller.hooks('before_client_request') assert call.method == 'SendStreamingMessage' - client_responses = setup.caller.hooks('after_agent_call') + client_responses = setup.caller.hooks('after_client_request') assert len(client_responses) == len(events) + 1 assert client_responses[-1].is_final # Server side: the header added by the client plugin arrives. - (run,) = setup.server.hooks('before_agent_run') + (run,) = setup.server.hooks('before_handler_run') assert run.method == 'SendStreamingMessage' assert incoming_headers(run.context)['x-trace'] == 'trace-123' assert run.agent_card.name == 'Plugin Agent' - server_responses = setup.server.hooks('after_agent_run') + server_responses = setup.server.hooks('after_handler_run') assert server_responses[-1].is_final assert sum(r.is_stream_event for r in server_responses) == len(events) @@ -273,10 +273,10 @@ async def test_unary_call_runs_hooks_on_both_sides(setup: Setup): task = await setup.client.get_task(GetTaskRequest(id=task_id)) assert task.id == task_id - assert [r.method for r in setup.caller.hooks('before_agent_call')] == [ + assert [r.method for r in setup.caller.hooks('before_client_request')] == [ 'GetTask' ] - (response,) = setup.server.hooks('after_agent_run') + (response,) = setup.server.hooks('after_handler_run') assert response.method == 'GetTask' assert response.payload.id == task_id assert response.is_final @@ -313,9 +313,9 @@ async def test_v03_jsonrpc_requests_run_server_hooks(): assert response.status_code == 200 assert response.json()['result']['status']['state'] == 'completed' - (run,) = recorder.hooks('before_agent_run') + (run,) = recorder.hooks('before_handler_run') assert run.method == 'SendMessage' - assert recorder.hooks('after_agent_run')[-1].is_final + assert recorder.hooks('after_handler_run')[-1].is_final _tenant: contextvars.ContextVar[str | None] = contextvars.ContextVar( @@ -332,7 +332,7 @@ async def before_execution(self, request_context: RequestContext) -> None: headers = request_context.call_context.state.get('headers', {}) _tenant.set(headers.get('x-tenant')) - async def before_agent_call(self, request: ClientRequest) -> None: + async def before_client_request(self, request: ClientRequest) -> None: if tenant := _tenant.get(): params = dict(request.context.service_parameters or {}) params['x-tenant'] = tenant @@ -401,5 +401,5 @@ async def cancel(self, context, event_queue): await downstream_http.aclose() assert response.status_code == 200, response.text - (run,) = downstream.hooks('before_agent_run') + (run,) = downstream.hooks('before_handler_run') assert run.context.state['headers']['x-tenant'] == 'acme' diff --git a/tests/plugins/test_manager.py b/tests/plugins/test_manager.py index d94c3a9ba..db0e479f9 100644 --- a/tests/plugins/test_manager.py +++ b/tests/plugins/test_manager.py @@ -38,17 +38,17 @@ def _record(self, hook: str) -> None: if hook in self.fail: raise ValueError(f'{self.name}.{hook} failed') - async def before_agent_run(self, request): - self._record('before_agent_run') + async def before_handler_run(self, request): + self._record('before_handler_run') - async def after_agent_run(self, response): - self._record('after_agent_run') + async def after_handler_run(self, response): + self._record('after_handler_run') - async def before_agent_call(self, request): - self._record('before_agent_call') + async def before_client_request(self, request): + self._record('before_client_request') - async def after_agent_call(self, response): - self._record('after_agent_call') + async def after_client_request(self, response): + self._record('after_client_request') async def before_execution(self, request_context): self._record('before_execution') @@ -93,26 +93,26 @@ async def test_before_in_order_after_in_reverse(): log: list[str] = [] m = plugins.sequential(Recorder('a', log), Recorder('b', log)) - await m.before_agent_run(server_request()) - await m.after_agent_run( + await m.before_handler_run(server_request()) + await m.after_handler_run( ServerResponse('GetTask', None, ServerCallContext()) ) - await m.before_agent_call(client_request()) - await m.after_agent_call( + await m.before_client_request(client_request()) + await m.after_client_request( ClientResponse('GetTask', None, ClientCallContext()) ) await m.before_execution(MagicMock()) await m.after_execution(MagicMock(), None) assert log == [ - 'a.before_agent_run', - 'b.before_agent_run', - 'b.after_agent_run', - 'a.after_agent_run', - 'a.before_agent_call', - 'b.before_agent_call', - 'b.after_agent_call', - 'a.after_agent_call', + 'a.before_handler_run', + 'b.before_handler_run', + 'b.after_handler_run', + 'a.after_handler_run', + 'a.before_client_request', + 'b.before_client_request', + 'b.after_client_request', + 'a.after_client_request', 'a.before_execution', 'b.before_execution', 'b.after_execution', @@ -154,14 +154,14 @@ async def test_observers_run_in_order(): ('before', 'after', 'call'), [ ( - 'before_agent_run', - 'after_agent_run', - lambda m: m.before_agent_run(server_request()), + 'before_handler_run', + 'after_handler_run', + lambda m: m.before_handler_run(server_request()), ), ( - 'before_agent_call', - 'after_agent_call', - lambda m: m.before_agent_call(client_request()), + 'before_client_request', + 'after_client_request', + lambda m: m.before_client_request(client_request()), ), ( 'before_execution', @@ -191,16 +191,16 @@ async def test_unwind_passes_error_to_after_hook(): seen: list[BaseException | None] = [] class Capture(ServerPlugin): - async def after_agent_run(self, response): + async def after_handler_run(self, response): seen.append(response.error) class Reject(ServerPlugin): - async def before_agent_run(self, request): + async def before_handler_run(self, request): raise PermissionError('denied') m = plugins.sequential(Capture(), Reject()) with pytest.raises(PermissionError): - await m.before_agent_run(server_request()) + await m.before_handler_run(server_request()) assert len(seen) == 1 assert isinstance(seen[0], PermissionError) @@ -209,8 +209,8 @@ async def before_agent_run(self, request): @pytest.mark.parametrize( 'hook', [ - 'after_agent_run', - 'after_agent_call', + 'after_handler_run', + 'after_client_request', 'after_execution', 'on_execution_event', 'on_task_event', @@ -222,10 +222,10 @@ async def test_after_and_observer_errors_are_swallowed(hook, caplog): log: list[str] = [] m = plugins.sequential(Recorder('a', log, fail={hook}), Recorder('b', log)) calls = { - 'after_agent_run': lambda: m.after_agent_run( + 'after_handler_run': lambda: m.after_handler_run( ServerResponse('GetTask', None, ServerCallContext()) ), - 'after_agent_call': lambda: m.after_agent_call( + 'after_client_request': lambda: m.after_client_request( ClientResponse('GetTask', None, ClientCallContext()) ), 'after_execution': lambda: m.after_execution(MagicMock(), None), @@ -293,18 +293,18 @@ async def test_nested_managers(): inner = plugins.sequential(Recorder('b', log), Recorder('c', log)) outer = plugins.sequential(Recorder('a', log), inner) - await outer.before_agent_run(server_request()) - await outer.after_agent_run( + await outer.before_handler_run(server_request()) + await outer.after_handler_run( ServerResponse('GetTask', None, ServerCallContext()) ) assert log == [ - 'a.before_agent_run', - 'b.before_agent_run', - 'c.before_agent_run', - 'c.after_agent_run', - 'b.after_agent_run', - 'a.after_agent_run', + 'a.before_handler_run', + 'b.before_handler_run', + 'c.before_handler_run', + 'c.after_handler_run', + 'b.after_handler_run', + 'a.after_handler_run', ] @@ -322,9 +322,9 @@ def test_rejects_non_plugins(): @pytest.mark.asyncio async def test_base_hooks_are_no_ops(): p = CompletePlugin() - assert await p.before_agent_run(server_request()) is None + assert await p.before_handler_run(server_request()) is None assert ( - await p.after_agent_call( + await p.after_client_request( ClientResponse('GetTask', None, ClientCallContext()) ) is None diff --git a/tests/server/request_handlers/test_default_request_handler_v2_plugins.py b/tests/server/request_handlers/test_default_request_handler_v2_plugins.py index 476d55375..17a6adf91 100644 --- a/tests/server/request_handlers/test_default_request_handler_v2_plugins.py +++ b/tests/server/request_handlers/test_default_request_handler_v2_plugins.py @@ -78,11 +78,11 @@ def hooks(self, name: str) -> list[Any]: def names(self) -> list[str]: return [hook for hook, _ in self.calls] - async def before_agent_run(self, request: ServerRequest) -> None: - self.calls.append(('before_agent_run', request)) + async def before_handler_run(self, request: ServerRequest) -> None: + self.calls.append(('before_handler_run', request)) - async def after_agent_run(self, response: ServerResponse) -> None: - self.calls.append(('after_agent_run', response)) + async def after_handler_run(self, response: ServerResponse) -> None: + self.calls.append(('after_handler_run', response)) async def before_execution(self, request_context: RequestContext) -> None: self.calls.append(('before_execution', request_context)) @@ -228,13 +228,13 @@ async def test_send_message_runs_all_server_hooks(): assert isinstance(result, Task) assert result.status.state == TaskState.TASK_STATE_COMPLETED - (request,) = recorder.hooks('before_agent_run') + (request,) = recorder.hooks('before_handler_run') assert request.method == 'SendMessage' assert request.context is context assert request.agent_card is CARD assert request.payload.message.message_id == 'm-1' - (response,) = recorder.hooks('after_agent_run') + (response,) = recorder.hooks('after_handler_run') assert response.method == 'SendMessage' assert response.payload is result assert response.error is None @@ -273,7 +273,7 @@ async def test_send_message_runs_all_server_hooks(): assert all(i.context is context for i in infos) names = recorder.names() - assert names[0] == 'before_agent_run' + assert names[0] == 'before_handler_run' assert names.index('before_execution') < names.index('on_execution_event') @@ -289,7 +289,7 @@ async def test_streaming_reports_each_event_and_a_final_call(): ) ] - responses = recorder.hooks('after_agent_run') + responses = recorder.hooks('after_handler_run') assert [r.payload for r in responses[:-1]] == streamed assert all(r.is_stream_event and not r.is_final for r in responses[:-1]) final = responses[-1] @@ -309,13 +309,13 @@ async def test_stream_closed_early_still_gets_final_call(): await anext(stream) await stream.aclose() - final = recorder.hooks('after_agent_run')[-1] + final = recorder.hooks('after_handler_run')[-1] assert final.is_final assert final.error is None @pytest.mark.asyncio -async def test_handler_error_reaches_after_agent_run(): +async def test_handler_error_reaches_after_handler_run(): recorder = Recorder() handler = make_handler(recorder) @@ -324,7 +324,7 @@ async def test_handler_error_reaches_after_agent_run(): GetTaskRequest(id='missing'), ServerCallContext() ) - (response,) = recorder.hooks('after_agent_run') + (response,) = recorder.hooks('after_handler_run') assert response.method == 'GetTask' assert isinstance(response.error, TaskNotFoundError) assert response.is_final @@ -349,14 +349,14 @@ async def test_stream_error_reaches_final_call(): ): pass - final = recorder.hooks('after_agent_run')[-1] + final = recorder.hooks('after_handler_run')[-1] assert isinstance(final.error, TaskNotFoundError) @pytest.mark.asyncio -async def test_before_agent_run_can_reject_the_request(): +async def test_before_handler_run_can_reject_the_request(): class Deny(plugins.ServerPlugin): - async def before_agent_run(self, request): + async def before_handler_run(self, request): raise PermissionError('denied') executed = False @@ -467,7 +467,7 @@ async def execute(self, context, event_queue): ) assert result.status.state == TaskState.TASK_STATE_SUBMITTED - assert recorder.hooks('after_agent_run')[0].is_final + assert recorder.hooks('after_handler_run')[0].is_final assert 'after_execution' not in recorder.names() release.set() From 13505d59d9d143c5a0ff6306714ab4a3dfa8c9c4 Mon Sep 17 00:00:00 2001 From: Mykyta Netipa Date: Thu, 22 Oct 2026 10:28:55 +0000 Subject: [PATCH 10/11] chore: allow PLW in spell check The plugin registry suppresses ruff's PLW0603 on its global statement, and check-spelling flags the rule prefix as an unknown word. --- .github/actions/spelling/allow.txt | 1 + 1 file changed, 1 insertion(+) diff --git a/.github/actions/spelling/allow.txt b/.github/actions/spelling/allow.txt index 0bd580314..2b03b4a23 100644 --- a/.github/actions/spelling/allow.txt +++ b/.github/actions/spelling/allow.txt @@ -115,6 +115,7 @@ opensource otherurl outerjoin pb2 +PLW podman Podman poolclass From 7f0c24209d27388d37b96b930a0a67b382ced799 Mon Sep 17 00:00:00 2001 From: Mykyta Netipa Date: Fri, 23 Oct 2026 16:48:11 +0000 Subject: [PATCH 11/11] refactor(plugins): use the client method names for plugin hooks Drop the rpc_method argument from BaseClient's call helpers: client plugins now receive the same method name as interceptors, and DefaultRequestHandlerV2 names its methods in @server_plugin_hooks with those same strings, so plugins see one naming scheme on both sides without a second parameter on every call. --- src/a2a/client/base_client.py | 33 +++--------- src/a2a/plugins/types.py | 9 ++-- .../server/request_handlers/_plugin_hooks.py | 2 +- .../default_request_handler_v2.py | 22 ++++---- tests/client/test_base_client_plugins.py | 53 ++++++------------- tests/integration/test_plugins.py | 10 ++-- ...test_default_request_handler_v2_plugins.py | 8 +-- 7 files changed, 49 insertions(+), 88 deletions(-) diff --git a/src/a2a/client/base_client.py b/src/a2a/client/base_client.py index b13873a52..fbcdb955e 100644 --- a/src/a2a/client/base_client.py +++ b/src/a2a/client/base_client.py @@ -91,7 +91,6 @@ 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 @@ -117,7 +116,6 @@ async def send_message( 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) @@ -179,7 +177,6 @@ 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 @@ -196,7 +193,6 @@ 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 @@ -221,7 +217,6 @@ 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 @@ -246,7 +241,6 @@ 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( @@ -273,7 +267,6 @@ 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( @@ -300,7 +293,6 @@ 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( @@ -324,7 +316,6 @@ 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( @@ -362,7 +353,6 @@ async def subscribe( 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 @@ -395,7 +385,6 @@ 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) @@ -414,11 +403,11 @@ async def close(self) -> None: def _client_request( self, input_data: Any, - rpc_method: str, + method: str, context: ClientCallContext | None, ) -> ClientRequest: return ClientRequest( - method=rpc_method, + method=method, payload=input_data, agent_card=self._card, context=context if context is not None else ClientCallContext(), @@ -432,13 +421,10 @@ 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``. + ``method`` is reported to both plugins and interceptors. """ plugins = self._plugins if not plugins.has_client_hooks: @@ -446,9 +432,7 @@ async def _execute_with_interceptors( input_data, method, context, transport_call ) - request = self._client_request( - input_data, rpc_method or method, context - ) + request = self._client_request(input_data, method, context) await plugins.before_client_request(request) try: result = await self._call_through_interceptors( @@ -514,13 +498,10 @@ async def _execute_stream_with_interceptors( transport_call: Callable[ [Any, ClientCallContext | None], 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``. + ``method`` is reported to both plugins and interceptors. """ plugins = self._plugins if not plugins.has_client_hooks: @@ -533,9 +514,7 @@ async def _execute_stream_with_interceptors( yield event return - request = self._client_request( - input_data, rpc_method or method, context - ) + request = self._client_request(input_data, method, context) await plugins.before_client_request(request) error: BaseException | None = None try: diff --git a/src/a2a/plugins/types.py b/src/a2a/plugins/types.py index 5fbb9f746..922dcdf0e 100644 --- a/src/a2a/plugins/types.py +++ b/src/a2a/plugins/types.py @@ -26,7 +26,7 @@ class ServerRequest: """An incoming request, passed to `ServerPlugin.before_handler_run`. Attributes: - method: A2A RPC name, e.g. ``'SendMessage'`` or ``'GetTask'``. + method: Method name, e.g. ``'send_message'`` or ``'get_task'``. payload: The request proto, e.g. ``SendMessageRequest``. context: The call context. For HTTP transports ``state['headers']`` holds the request headers. @@ -49,7 +49,7 @@ class ServerResponse: ends, including on error or client disconnect. Attributes: - method: A2A RPC name. + method: Method name. payload: The handler result or stream event, ``None`` on error or stream end. context: The same call context passed to `before_handler_run`. @@ -71,7 +71,8 @@ class ClientRequest: """An outgoing call, passed to `ClientPlugin.before_client_request`. Attributes: - method: A2A RPC name, matching the server side. + method: Method name, the same as on the server side and as passed + to interceptors. payload: The request proto. Plugins may replace it. agent_card: The card of the remote agent. context: The call context. ``service_parameters`` holds the outgoing @@ -92,7 +93,7 @@ class ClientResponse: `ServerResponse`. Attributes: - method: A2A RPC name. + method: Method name. payload: The response or stream event, ``None`` on error or stream end. context: The same call context passed to `before_client_request`. error: The exception raised by the call, if any. diff --git a/src/a2a/server/request_handlers/_plugin_hooks.py b/src/a2a/server/request_handlers/_plugin_hooks.py index adf0c02a3..4a0c99591 100644 --- a/src/a2a/server/request_handlers/_plugin_hooks.py +++ b/src/a2a/server/request_handlers/_plugin_hooks.py @@ -40,7 +40,7 @@ def server_plugin_hooks(method: str) -> Callable[[_F], _F]: has no server hooks the method runs unchanged. Args: - method: The A2A RPC name passed to the hooks, e.g. ``'GetTask'``. + method: The method name passed to the hooks, e.g. ``'get_task'``. """ def decorator(func: _F) -> _F: diff --git a/src/a2a/server/request_handlers/default_request_handler_v2.py b/src/a2a/server/request_handlers/default_request_handler_v2.py index 16422da49..4121ae80f 100644 --- a/src/a2a/server/request_handlers/default_request_handler_v2.py +++ b/src/a2a/server/request_handlers/default_request_handler_v2.py @@ -221,7 +221,7 @@ async def aclose(self) -> None: """ await self._active_task_registry.aclose() - @server_plugin_hooks('GetTask') + @server_plugin_hooks('get_task') @validate_request_params async def on_get_task( # noqa: D102 self, @@ -237,7 +237,7 @@ async def on_get_task( # noqa: D102 return apply_history_length(stored.task, params) - @server_plugin_hooks('ListTasks') + @server_plugin_hooks('list_tasks') @validate_request_params async def on_list_tasks( # noqa: D102 self, @@ -259,7 +259,7 @@ async def on_list_tasks( # noqa: D102 return page - @server_plugin_hooks('CancelTask') + @server_plugin_hooks('cancel_task') @validate_request_params async def on_cancel_task( # noqa: D102 self, @@ -437,7 +437,7 @@ async def _setup_active_task( return active_task, request_context - @server_plugin_hooks('SendMessage') + @server_plugin_hooks('send_message') @validate_request_params async def on_message_send( # noqa: D102 self, @@ -503,7 +503,7 @@ async def on_message_send( # noqa: D102 ) return result - @server_plugin_hooks('SendStreamingMessage') + @server_plugin_hooks('send_message_streaming') @validate_request_params @validate( lambda self: self._agent_card.capabilities.streaming, @@ -536,7 +536,7 @@ async def on_message_send_stream( # noqa: D102 else: yield event - @server_plugin_hooks('CreateTaskPushNotificationConfig') + @server_plugin_hooks('create_task_push_notification_config') @validate_request_params @validate( lambda self: self._agent_card.capabilities.push_notifications, @@ -567,7 +567,7 @@ async def on_create_task_push_notification_config( # noqa: D102 # Custom stores written before set_info returned the stored config. return normalize_push_notification_config(task_id, params) - @server_plugin_hooks('GetTaskPushNotificationConfig') + @server_plugin_hooks('get_task_push_notification_config') @validate_request_params @validate( lambda self: self._agent_card.capabilities.push_notifications, @@ -597,7 +597,7 @@ async def on_get_task_push_notification_config( # noqa: D102 raise TaskNotFoundError - @server_plugin_hooks('SubscribeToTask') + @server_plugin_hooks('subscribe') @validate_request_params @validate( lambda self: self._agent_card.capabilities.streaming, @@ -672,7 +672,7 @@ async def _subscribe_remote( ): return - @server_plugin_hooks('ListTaskPushNotificationConfigs') + @server_plugin_hooks('list_task_push_notification_configs') @validate_request_params @validate( lambda self: self._agent_card.capabilities.push_notifications, @@ -699,7 +699,7 @@ async def on_list_task_push_notification_configs( # noqa: D102 configs=push_notification_config_list ) - @server_plugin_hooks('DeleteTaskPushNotificationConfig') + @server_plugin_hooks('delete_task_push_notification_config') @validate_request_params @validate( lambda self: self._agent_card.capabilities.push_notifications, @@ -721,7 +721,7 @@ async def on_delete_task_push_notification_config( # noqa: D102 await self._push_config_store.delete_info(task_id, context, config_id) - @server_plugin_hooks('GetExtendedAgentCard') + @server_plugin_hooks('get_extended_agent_card') @validate_request_params @validate( lambda self: self._agent_card.capabilities.extended_agent_card, diff --git a/tests/client/test_base_client_plugins.py b/tests/client/test_base_client_plugins.py index 8a1ed4e09..2704f9291 100644 --- a/tests/client/test_base_client_plugins.py +++ b/tests/client/test_base_client_plugins.py @@ -115,7 +115,7 @@ async def test_unary_call_runs_hooks(): assert result.id == 't-1' (request,) = recorder.hooks('before_client_request') - assert request.method == 'GetTask' + assert request.method == 'get_task' assert request.agent_card.name == 'Remote' (response,) = recorder.hooks('after_client_request') assert response.payload is result @@ -142,7 +142,7 @@ async def test_non_streaming_send_message_is_a_unary_call(): ] assert len(events) == 1 - assert recorder.hooks('before_client_request')[0].method == 'SendMessage' + assert recorder.hooks('before_client_request')[0].method == 'send_message' (response,) = recorder.hooks('after_client_request') assert response.is_final @@ -169,7 +169,7 @@ async def stream(request, *, context=None): assert events == streamed (request,) = recorder.hooks('before_client_request') - assert request.method == 'SendStreamingMessage' + assert request.method == 'send_message_streaming' responses = recorder.hooks('after_client_request') assert [r.payload for r in responses[:-1]] == streamed assert all(r.is_stream_event and not r.is_final for r in responses[:-1]) @@ -321,54 +321,54 @@ def test_factory_clients_use_the_global_manager(): @pytest.mark.parametrize( - ('call', 'rpc_name'), + ('call', 'method'), [ - (lambda c: c.list_tasks(ListTasksRequest()), 'ListTasks'), - (lambda c: c.cancel_task(CancelTaskRequest(id='t-1')), 'CancelTask'), + (lambda c: c.list_tasks(ListTasksRequest()), 'list_tasks'), + (lambda c: c.cancel_task(CancelTaskRequest(id='t-1')), 'cancel_task'), ( lambda c: c.create_task_push_notification_config( TaskPushNotificationConfig(task_id='t-1', url='http://hook') ), - 'CreateTaskPushNotificationConfig', + 'create_task_push_notification_config', ), ( lambda c: c.get_task_push_notification_config( GetTaskPushNotificationConfigRequest(task_id='t-1', id='c-1') ), - 'GetTaskPushNotificationConfig', + 'get_task_push_notification_config', ), ( lambda c: c.list_task_push_notification_configs( ListTaskPushNotificationConfigsRequest(task_id='t-1') ), - 'ListTaskPushNotificationConfigs', + 'list_task_push_notification_configs', ), ( lambda c: c.delete_task_push_notification_config( DeleteTaskPushNotificationConfigRequest(task_id='t-1', id='c-1') ), - 'DeleteTaskPushNotificationConfig', + 'delete_task_push_notification_config', ), ( lambda c: c.get_extended_agent_card(GetExtendedAgentCardRequest()), - 'GetExtendedAgentCard', + 'get_extended_agent_card', ), ], ) @pytest.mark.asyncio -async def test_unary_methods_report_rpc_names(call, rpc_name): +async def test_unary_methods_report_method_names(call, method): recorder = Recorder() await call(client(recorder, AsyncMock(spec=ClientTransport))) assert [r.method for r in recorder.hooks('before_client_request')] == [ - rpc_name + method ] - assert recorder.hooks('after_client_request')[-1].method == rpc_name + assert recorder.hooks('after_client_request')[-1].method == method @pytest.mark.asyncio -async def test_subscribe_reports_rpc_name(): +async def test_subscribe_reports_method_name(): recorder = Recorder() transport = AsyncMock(spec=ClientTransport) @@ -382,26 +382,7 @@ async def stream(request, *, context=None): ): pass - assert ( - recorder.hooks('before_client_request')[0].method == 'SubscribeToTask' - ) + assert recorder.hooks('before_client_request')[0].method == 'subscribe' assert all( - r.method == 'SubscribeToTask' - for r in recorder.hooks('after_client_request') + r.method == 'subscribe' for r in recorder.hooks('after_client_request') ) - - -@pytest.mark.asyncio -async def test_rpc_method_defaults_to_method(): - recorder = Recorder() - transport = AsyncMock(spec=ClientTransport) - c = client(recorder, transport) - - await c._execute_with_interceptors( - input_data=GetTaskRequest(id='t-1'), - method='custom_call', - context=None, - transport_call=AsyncMock(return_value=Task(id='t-1')), - ) - - assert recorder.hooks('before_client_request')[0].method == 'custom_call' diff --git a/tests/integration/test_plugins.py b/tests/integration/test_plugins.py index 16a9d1056..c5f1d47a7 100644 --- a/tests/integration/test_plugins.py +++ b/tests/integration/test_plugins.py @@ -240,14 +240,14 @@ async def test_streaming_message_runs_hooks_on_both_sides(setup: Setup): # Client side. (call,) = setup.caller.hooks('before_client_request') - assert call.method == 'SendStreamingMessage' + assert call.method == 'send_message_streaming' client_responses = setup.caller.hooks('after_client_request') assert len(client_responses) == len(events) + 1 assert client_responses[-1].is_final # Server side: the header added by the client plugin arrives. (run,) = setup.server.hooks('before_handler_run') - assert run.method == 'SendStreamingMessage' + assert run.method == 'send_message_streaming' assert incoming_headers(run.context)['x-trace'] == 'trace-123' assert run.agent_card.name == 'Plugin Agent' server_responses = setup.server.hooks('after_handler_run') @@ -274,10 +274,10 @@ async def test_unary_call_runs_hooks_on_both_sides(setup: Setup): assert task.id == task_id assert [r.method for r in setup.caller.hooks('before_client_request')] == [ - 'GetTask' + 'get_task' ] (response,) = setup.server.hooks('after_handler_run') - assert response.method == 'GetTask' + assert response.method == 'get_task' assert response.payload.id == task_id assert response.is_final @@ -314,7 +314,7 @@ async def test_v03_jsonrpc_requests_run_server_hooks(): assert response.status_code == 200 assert response.json()['result']['status']['state'] == 'completed' (run,) = recorder.hooks('before_handler_run') - assert run.method == 'SendMessage' + assert run.method == 'send_message' assert recorder.hooks('after_handler_run')[-1].is_final diff --git a/tests/server/request_handlers/test_default_request_handler_v2_plugins.py b/tests/server/request_handlers/test_default_request_handler_v2_plugins.py index 17a6adf91..222607110 100644 --- a/tests/server/request_handlers/test_default_request_handler_v2_plugins.py +++ b/tests/server/request_handlers/test_default_request_handler_v2_plugins.py @@ -229,13 +229,13 @@ async def test_send_message_runs_all_server_hooks(): assert result.status.state == TaskState.TASK_STATE_COMPLETED (request,) = recorder.hooks('before_handler_run') - assert request.method == 'SendMessage' + assert request.method == 'send_message' assert request.context is context assert request.agent_card is CARD assert request.payload.message.message_id == 'm-1' (response,) = recorder.hooks('after_handler_run') - assert response.method == 'SendMessage' + assert response.method == 'send_message' assert response.payload is result assert response.error is None assert response.is_final @@ -293,7 +293,7 @@ async def test_streaming_reports_each_event_and_a_final_call(): assert [r.payload for r in responses[:-1]] == streamed assert all(r.is_stream_event and not r.is_final for r in responses[:-1]) final = responses[-1] - assert final.method == 'SendStreamingMessage' + assert final.method == 'send_message_streaming' assert final.is_final assert not final.is_stream_event assert final.payload is None @@ -325,7 +325,7 @@ async def test_handler_error_reaches_after_handler_run(): ) (response,) = recorder.hooks('after_handler_run') - assert response.method == 'GetTask' + assert response.method == 'get_task' assert isinstance(response.error, TaskNotFoundError) assert response.is_final