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 diff --git a/src/a2a/client/base_client.py b/src/a2a/client/base_client.py index 763f23fb5..fbcdb955e 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, @@ -32,7 +35,11 @@ class BaseClient(Client): - """Base implementation of the A2A client, containing transport-independent logic.""" + """Base implementation of the A2A client, containing transport-independent logic. + + Client plugin hooks run around every call: ``before_client_request`` before + the interceptors, ``after_client_request`` after them. + """ def __init__( self, @@ -40,12 +47,25 @@ def __init__( config: ClientConfig, transport: ClientTransport, interceptors: list[ClientCallInterceptor], + *, + plugins: PluginManager | None = None, ): + """Initializes the client. + + Args: + card: The card of the agent to call. + config: The client configuration. + transport: The transport used for calls. + interceptors: Interceptors applied to every call. + plugins: Plugin manager whose client hooks run for every call. + Defaults to the global manager at construction time. + """ super().__init__(interceptors) self._card = card self._config = config self._transport = transport self._interceptors = interceptors + self._plugins = plugins if plugins is not None else get_global() async def send_message( self, @@ -90,15 +110,20 @@ async def send_message( yield stream_response return - async for event in self._execute_stream_with_interceptors( - input_data=request, - method='send_message_streaming', - context=context, - transport_call=lambda req, ctx: ( - self._transport.send_message_streaming(req, context=ctx) - ), - ): - yield event + # aclosing() closes the inner stream (and the transport stream) as + # soon as the caller stops iterating. + async with aclosing( + self._execute_stream_with_interceptors( + input_data=request, + method='send_message_streaming', + context=context, + transport_call=lambda req, ctx: ( + self._transport.send_message_streaming(req, context=ctx) + ), + ) + ) as events: + async for event in events: + yield event def _apply_client_config(self, request: SendMessageRequest) -> None: request.configuration.return_immediately |= self._config.polling @@ -324,15 +349,18 @@ async def subscribe( 'client and/or server do not support resubscription.' ) - async for event in self._execute_stream_with_interceptors( - input_data=request, - method='subscribe', - context=context, - transport_call=lambda req, ctx: self._transport.subscribe( - req, context=ctx - ), - ): - yield event + async with aclosing( + self._execute_stream_with_interceptors( + input_data=request, + method='subscribe', + context=context, + transport_call=lambda req, ctx: self._transport.subscribe( + req, context=ctx + ), + ) + ) as events: + async for event in events: + yield event async def get_extended_agent_card( self, @@ -372,6 +400,19 @@ async def close(self) -> None: """Closes the underlying transport.""" await self._transport.close() + def _client_request( + self, + input_data: Any, + method: str, + context: ClientCallContext | None, + ) -> ClientRequest: + return ClientRequest( + method=method, + payload=input_data, + agent_card=self._card, + context=context if context is not None else ClientCallContext(), + ) + async def _execute_with_interceptors( self, input_data: Any, @@ -380,6 +421,41 @@ async def _execute_with_interceptors( transport_call: Callable[ [Any, ClientCallContext | None], Awaitable[Any] ], + ) -> Any: + """Runs a unary call through plugins, interceptors and the transport. + + ``method`` is reported to both plugins and interceptors. + """ + plugins = self._plugins + if not plugins.has_client_hooks: + return await self._call_through_interceptors( + input_data, method, context, transport_call + ) + + request = self._client_request(input_data, method, context) + await plugins.before_client_request(request) + try: + result = await self._call_through_interceptors( + request.payload, method, request.context, transport_call + ) + except Exception as e: + await plugins.after_client_request( + ClientResponse(request.method, None, request.context, error=e) + ) + raise + await plugins.after_client_request( + ClientResponse(request.method, result, request.context) + ) + return result + + async def _call_through_interceptors( + self, + input_data: Any, + method: str, + context: ClientCallContext | None, + transport_call: Callable[ + [Any, ClientCallContext | None], Awaitable[Any] + ], ) -> Any: before_args = BeforeArgs( input=input_data, @@ -422,8 +498,62 @@ async def _execute_stream_with_interceptors( transport_call: Callable[ [Any, ClientCallContext | None], AsyncIterator[StreamResponse] ], - ) -> AsyncIterator[StreamResponse]: + ) -> AsyncGenerator[StreamResponse, None]: + """Runs a streaming call through plugins, interceptors and the transport. + ``method`` is reported to both plugins and interceptors. + """ + plugins = self._plugins + if not plugins.has_client_hooks: + async with aclosing( + self._stream_through_interceptors( + input_data, method, context, transport_call + ) + ) as events: + async for event in events: + yield event + return + + request = self._client_request(input_data, method, context) + await plugins.before_client_request(request) + error: BaseException | None = None + try: + async with aclosing( + self._stream_through_interceptors( + request.payload, method, request.context, transport_call + ) + ) as events: + async for event in events: + await plugins.after_client_request( + ClientResponse( + request.method, + event, + request.context, + is_stream_event=True, + is_final=False, + ) + ) + yield event + except Exception as e: + error = e + raise + finally: + # Also runs when the caller stops iterating early. + await plugins.after_client_request( + ClientResponse( + request.method, None, request.context, error=error + ) + ) + + async def _stream_through_interceptors( + self, + input_data: Any, + method: str, + context: ClientCallContext | None, + transport_call: Callable[ + [Any, ClientCallContext | None], AsyncIterator[StreamResponse] + ], + ) -> AsyncGenerator[StreamResponse, None]: before_args = BeforeArgs( input=input_data, method=method, 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/base.py b/src/a2a/plugins/base.py new file mode 100644 index 000000000..23a51eb53 --- /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_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_handler_run(self, response: ServerResponse) -> None: + """Called with the handler result, per stream event, and at the end.""" + + +class ClientPlugin(Plugin): + """Hooks around each request made by an A2A `Client`.""" + + async def before_client_request(self, request: ClientRequest) -> None: + """Called before the request is sent. Raise to cancel it.""" + + async def after_client_request(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..0b2af80d9 --- /dev/null +++ b/src/a2a/plugins/manager.py @@ -0,0 +1,296 @@ +"""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_handler_run(self, request: ServerRequest) -> None: # noqa: D102 + ran: list[ServerPlugin] = [] + for plugin in self._server: + try: + 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_handler_run', + done.after_handler_run(response), + ) + raise + ran.append(plugin) + + async def after_handler_run(self, response: ServerResponse) -> None: # noqa: D102 + for plugin in reversed(self._server): + await _run_safely( + plugin, 'after_handler_run', plugin.after_handler_run(response) + ) + + # --- client --- + + async def before_client_request(self, request: ClientRequest) -> None: # noqa: D102 + ran: list[ClientPlugin] = [] + for plugin in self._client: + try: + await plugin.before_client_request(request) + except Exception as e: + response = ClientResponse( + request.method, None, request.context, error=e + ) + for done in reversed(ran): + await _run_safely( + done, + 'after_client_request', + done.after_client_request(response), + ) + raise + ran.append(plugin) + + async def after_client_request(self, response: ClientResponse) -> None: # noqa: D102 + for plugin in reversed(self._client): + await _run_safely( + plugin, + 'after_client_request', + plugin.after_client_request(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, clients and push senders created afterwards. + + `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): + 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..922dcdf0e --- /dev/null +++ b/src/a2a/plugins/types.py @@ -0,0 +1,174 @@ +"""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_handler_run`. + + Attributes: + 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. + 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_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 + once more with ``is_final=True`` and ``payload=None`` when the stream + ends, including on error or client disconnect. + + Attributes: + 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`. + 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_client_request`. + + Attributes: + 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 + headers. + """ + + method: str + payload: Any + agent_card: AgentCard + context: ClientCallContext + + +@dataclass +class ClientResponse: + """A call result, passed to `ClientPlugin.after_client_request`. + + Follows the same ``is_stream_event`` / ``is_final`` rules as + `ServerResponse`. + + Attributes: + 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. + 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/src/a2a/server/agent_execution/active_task_registry.py b/src/a2a/server/agent_execution/active_task_registry.py index e5daaaa2e..a5350bbf4 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 @@ -28,15 +29,17 @@ class ActiveTaskRegistry: """A registry for active ActiveTask instances.""" - def __init__( + def __init__( # noqa: PLR0913 self, agent_executor: AgentExecutor, task_store: TaskStore | VersionedTaskStore, push_sender: PushNotificationSender | None = None, event_stream: TaskEventStream | None = None, interrupted_task_timeout: float | 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._interrupted_task_timeout = interrupted_task_timeout @@ -72,6 +75,7 @@ async def get_or_create( task_store=self._task_store, initial_message=initial_message, context=call_context, + plugins=self._plugins, ) active_task = ActiveTask( 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/request_handlers/_plugin_hooks.py b/src/a2a/server/request_handlers/_plugin_hooks.py new file mode 100644 index 000000000..4a0c99591 --- /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_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_handler_run`. When the handler's manager + has no server hooks the method runs unchanged. + + Args: + method: The method name passed to the hooks, e.g. ``'get_task'``. + """ + + 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_handler_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_handler_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_handler_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_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_handler_run( + ServerResponse(method, None, context, error=e) + ) + raise + await plugins.after_handler_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 ee32880e7..895185282 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 @@ -101,6 +108,10 @@ class DefaultRequestHandlerV2(RequestHandler): ``input-required`` or ``auth-required`` after that many seconds; a follow-up recreates it from the store. ``None`` (default) never releases; any other value must be greater than 0. + + ``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] @@ -124,6 +135,7 @@ def __init__( # noqa: PLR0913 validate_input_modes: bool = False, event_stream: TaskEventStream | None = None, interrupted_task_timeout: float | None = None, + plugins: PluginManager | None = None, ) -> None: if ( interrupted_task_timeout is not None @@ -141,6 +153,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 = ( @@ -179,14 +194,29 @@ 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, interrupted_task_timeout=interrupted_task_timeout, + 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: @@ -203,6 +233,7 @@ async def aclose(self) -> None: """ await self._active_task_registry.aclose() + @server_plugin_hooks('get_task') @validate_request_params async def on_get_task( # noqa: D102 self, @@ -218,6 +249,7 @@ async def on_get_task( # noqa: D102 return apply_history_length(stored.task, params) + @server_plugin_hooks('list_tasks') @validate_request_params async def on_list_tasks( # noqa: D102 self, @@ -239,6 +271,7 @@ async def on_list_tasks( # noqa: D102 return page + @server_plugin_hooks('cancel_task') @validate_request_params async def on_cancel_task( # noqa: D102 self, @@ -329,6 +362,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) @@ -408,6 +453,7 @@ async def _setup_active_task( return active_task, request_context + @server_plugin_hooks('send_message') @validate_request_params async def on_message_send( # noqa: D102 self, @@ -473,6 +519,7 @@ async def on_message_send( # noqa: D102 ) return result + @server_plugin_hooks('send_message_streaming') @validate_request_params @validate( lambda self: self._agent_card.capabilities.streaming, @@ -505,6 +552,7 @@ async def on_message_send_stream( # noqa: D102 else: yield event + @server_plugin_hooks('create_task_push_notification_config') @validate_request_params @validate( lambda self: self._agent_card.capabilities.push_notifications, @@ -535,6 +583,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('get_task_push_notification_config') @validate_request_params @validate( lambda self: self._agent_card.capabilities.push_notifications, @@ -564,6 +613,7 @@ async def on_get_task_push_notification_config( # noqa: D102 raise TaskNotFoundError + @server_plugin_hooks('subscribe') @validate_request_params @validate( lambda self: self._agent_card.capabilities.streaming, @@ -638,6 +688,7 @@ async def _subscribe_remote( ): return + @server_plugin_hooks('list_task_push_notification_configs') @validate_request_params @validate( lambda self: self._agent_card.capabilities.push_notifications, @@ -664,6 +715,7 @@ async def on_list_task_push_notification_configs( # noqa: D102 configs=push_notification_config_list ) + @server_plugin_hooks('delete_task_push_notification_config') @validate_request_params @validate( lambda self: self._agent_card.capabilities.push_notifications, @@ -685,6 +737,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('get_extended_agent_card') @validate_request_params @validate( lambda self: self._agent_card.capabilities.extended_agent_card, diff --git a/src/a2a/server/tasks/base_push_notification_sender.py b/src/a2a/server/tasks/base_push_notification_sender.py index 070361a52..bade95752 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,8 @@ from google.protobuf.json_format import MessageToDict +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 ( PushNotificationConfigStore, @@ -32,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. @@ -50,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( @@ -64,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 @@ -121,39 +130,93 @@ async def _dispatch_notification( return False if not accepted: return False - 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, - ) + plugins = self._plugins if self._plugins.has_push_hooks else None + request = PushRequest( + task_id, url, self._build_headers(push_info, task_id), event + ) + + 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 + error: Exception | None = None + try: response = await self._client.post( url, json=MessageToDict(to_stream_response(event)), - headers=headers, + headers=request.headers, ) + status_code = response.status_code response.raise_for_status() - logger.info( - 'Push-notification sent for task_id=%s to URL: %s', task_id, url - ) - except Exception: + except Exception as e: + error = e logger.exception( 'Error sending push-notification for task_id=%s to URL: %s.', task_id, url, ) - return False - return True + else: + logger.info( + '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') 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/client/test_base_client_plugins.py b/tests/client/test_base_client_plugins.py new file mode 100644 index 000000000..2704f9291 --- /dev/null +++ b/tests/client/test_base_client_plugins.py @@ -0,0 +1,388 @@ +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, + CancelTaskRequest, + DeleteTaskPushNotificationConfigRequest, + GetExtendedAgentCardRequest, + GetTaskPushNotificationConfigRequest, + GetTaskRequest, + ListTaskPushNotificationConfigsRequest, + ListTasksRequest, + Message, + Part, + Role, + SendMessageRequest, + SendMessageResponse, + StreamResponse, + SubscribeToTaskRequest, + Task, + TaskPushNotificationConfig, + 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_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_client_request', request)) + + async def after_client_request(self, response: ClientResponse) -> None: + self.calls.append(('after_client_request', 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, + ) + + +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(): + 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_client_request') + assert request.method == 'get_task' + assert request.agent_card.name == 'Remote' + (response,) = recorder.hooks('after_client_request') + 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_client_request')[0].method == 'send_message' + (response,) = recorder.hooks('after_client_request') + 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_client_request') + 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]) + 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_client_request')[-1].is_final + + +@pytest.mark.asyncio +async def test_transport_errors_reach_after_client_request(): + 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_client_request') + 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_client_request')[-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_client_request(self, request): + order.append('plugin.before') + + async def after_client_request(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_client_request_can_cancel_the_call(): + class Deny(plugins.ClientPlugin): + async def before_client_request(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_client_request')[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_clients_use_the_global_manager(): + 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 + + +@pytest.mark.parametrize( + ('call', 'method'), + [ + (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') + ), + 'create_task_push_notification_config', + ), + ( + lambda c: c.get_task_push_notification_config( + GetTaskPushNotificationConfigRequest(task_id='t-1', id='c-1') + ), + 'get_task_push_notification_config', + ), + ( + lambda c: c.list_task_push_notification_configs( + ListTaskPushNotificationConfigsRequest(task_id='t-1') + ), + 'list_task_push_notification_configs', + ), + ( + lambda c: c.delete_task_push_notification_config( + DeleteTaskPushNotificationConfigRequest(task_id='t-1', id='c-1') + ), + 'delete_task_push_notification_config', + ), + ( + lambda c: c.get_extended_agent_card(GetExtendedAgentCardRequest()), + 'get_extended_agent_card', + ), + ], +) +@pytest.mark.asyncio +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')] == [ + method + ] + assert recorder.hooks('after_client_request')[-1].method == method + + +@pytest.mark.asyncio +async def test_subscribe_reports_method_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_client_request')[0].method == 'subscribe' + assert all( + r.method == 'subscribe' for r in recorder.hooks('after_client_request') + ) diff --git a/tests/integration/test_plugins.py b/tests/integration/test_plugins.py new file mode 100644 index 000000000..c5f1d47a7 --- /dev/null +++ b/tests/integration/test_plugins.py @@ -0,0 +1,405 @@ +"""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.context import ServerCallContext +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_handler_run(self, request: ServerRequest) -> None: + self.calls.append(('before_handler_run', request)) + + async def after_handler_run(self, response: ServerResponse) -> None: + self.calls.append(('after_handler_run', response)) + + 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_client_request', request)) + + 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)) + + 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)] + ) + ) + + +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(): + 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_client_request') + 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 == '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') + 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_client_request')] == [ + 'get_task' + ] + (response,) = setup.server.hooks('after_handler_run') + assert response.method == 'get_task' + 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_handler_run') + assert run.method == 'send_message' + assert recorder.hooks('after_handler_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_client_request(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_handler_run') + assert run.context.state['headers']['x-tenant'] == 'acme' diff --git a/tests/plugins/test_manager.py b/tests/plugins/test_manager.py new file mode 100644 index 000000000..db0e479f9 --- /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_handler_run(self, request): + self._record('before_handler_run') + + async def after_handler_run(self, response): + self._record('after_handler_run') + + async def before_client_request(self, request): + self._record('before_client_request') + + async def after_client_request(self, response): + self._record('after_client_request') + + 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_handler_run(server_request()) + await m.after_handler_run( + ServerResponse('GetTask', None, ServerCallContext()) + ) + 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_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', + '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_handler_run', + 'after_handler_run', + lambda m: m.before_handler_run(server_request()), + ), + ( + 'before_client_request', + 'after_client_request', + lambda m: m.before_client_request(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_handler_run(self, response): + seen.append(response.error) + + class Reject(ServerPlugin): + async def before_handler_run(self, request): + raise PermissionError('denied') + + m = plugins.sequential(Capture(), Reject()) + with pytest.raises(PermissionError): + await m.before_handler_run(server_request()) + + assert len(seen) == 1 + assert isinstance(seen[0], PermissionError) + + +@pytest.mark.parametrize( + 'hook', + [ + 'after_handler_run', + 'after_client_request', + '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_handler_run': lambda: m.after_handler_run( + ServerResponse('GetTask', None, ServerCallContext()) + ), + 'after_client_request': lambda: m.after_client_request( + 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_handler_run(server_request()) + await outer.after_handler_run( + ServerResponse('GetTask', None, ServerCallContext()) + ) + + assert log == [ + 'a.before_handler_run', + 'b.before_handler_run', + 'c.before_handler_run', + 'c.after_handler_run', + 'b.after_handler_run', + 'a.after_handler_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_handler_run(server_request()) is None + assert ( + await p.after_client_request( + 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] 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..222607110 --- /dev/null +++ b/tests/server/request_handlers/test_default_request_handler_v2_plugins.py @@ -0,0 +1,707 @@ +"""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_handler_run(self, request: ServerRequest) -> None: + self.calls.append(('before_handler_run', request)) + + 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)) + + 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_handler_run') + 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 == 'send_message' + 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_handler_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_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] + assert final.method == 'send_message_streaming' + 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_handler_run')[-1] + assert final.is_final + assert final.error is None + + +@pytest.mark.asyncio +async def test_handler_error_reaches_after_handler_run(): + recorder = Recorder() + handler = make_handler(recorder) + + with pytest.raises(TaskNotFoundError): + await handler.on_get_task( + GetTaskRequest(id='missing'), ServerCallContext() + ) + + (response,) = recorder.hooks('after_handler_run') + assert response.method == 'get_task' + 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_handler_run')[-1] + assert isinstance(final.error, TaskNotFoundError) + + +@pytest.mark.asyncio +async def test_before_handler_run_can_reject_the_request(): + class Deny(plugins.ServerPlugin): + async def before_handler_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_handler_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], + manager: plugins.PluginManager | None = None, +) -> 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, + plugins=manager, + ) + 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(): + 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 + ) + + 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( + {}, plugins.sequential(plugin, Seen()) + ) + handler = DefaultRequestHandlerV2( + agent_executor=TaskAgent(), + task_store=InMemoryTaskStore(), + agent_card=CARD, + push_config_store=config_store, + push_sender=sender, + ) + 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] = [] + recorder = Recorder() + 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, + plugins=plugins.sequential(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_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({}) + 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 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.hooks('before_push_send') == [] 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))