Repository navigation
feat: add plugin framework - #1352
mykytanetipa wants to merge 13 commits into
Conversation
…Manager Implements the plugin model from go/a2a-plugins for the Python SDK: a plugin is one object that implements any of ServerPlugin, ClientPlugin, ExecutionPlugin, TaskObserverPlugin and PushPlugin, so cross-cutting concerns such as telemetry or auth plug in with one registration call instead of per-SDK interceptor wiring. SequentialPluginManager runs before-hooks in registration order and after-hooks in reverse, unwinds plugins that already ran when a before-hook rejects a call, and logs and swallows after-hook and observer failures so a buggy plugin cannot break a request. Its has_* properties let call sites skip all plugin work when nothing is registered.
DefaultRequestHandlerV2 takes a plugins manager (defaulting to the global one) and runs: - before/after_agent_run around every public on_* method, with a call per stream event and a final call on completion, error or disconnect; - before/after_execution and on_execution_event around the agent's execute(), via a wrapping executor and event queue, so plugins see the detached execution and the time of each agent event; - on_task_event after every successful task write. TaskManager tracks the persisted state itself because it mutates the cached task in place before saving, which made the store's prev argument useless for from_state; - before_push_send/on_push_sent per webhook delivery, scoped to the V2 path through a contextvar so the shared push sender stays hook-free for LegacyRequestHandler. Transports record the protocol in ServerCallContext.state['a2a.transport'] and gRPC exposes invocation metadata as state['headers'], so plugins can read trace and tenant headers on every transport. With no plugins registered nothing is wrapped and behavior is unchanged.
Every client call now runs before_agent_call before the interceptors and after_agent_call after them, with one call per stream event, a final call when the stream ends, and the error when the transport fails. The hooks get the RPC name used on the server side and the protocol binding and URL ClientFactory selected, so a plugin can tag calls and add headers such as trace context. send_message and subscribe now close the inner stream with aclosing() so an early break by the caller releases the transport stream and runs the final hook right away instead of at garbage collection.
Runs client and server hooks against in-process servers on every transport, checks that a header added by a client plugin reaches the server, that v0.3 JSON-RPC requests reach the V2 hooks, and that state set in before_execution propagates to calls the agent makes to a second agent.
Replace the separate _post_with_plugins method with a single delivery path in _dispatch_notification. Header building moves to _build_headers and the on_push_sent call to _report, which both paths share. Without a push plugin manager the only added cost per delivery is constructing a PushRequest and reading a monotonic clock; behavior is unchanged.
The plugin hooks do not depend on these additions; they were added ahead of a telemetry plugin that is not part of this change and can be added back compatibly when it lands: - ClientRequest.transport/url and the ClientFactory wiring that filled them; - ServerCallContext.state['a2a.transport'] written by the JSON-RPC, REST, gRPC and v0.3 dispatchers; - gRPC invocation metadata copied into state['headers']. Plugins can still read it through state['grpc_context']. The dispatchers, v0.3 adapters, gRPC handler, ClientFactory and ServerCallContext are back to their state on main.
Replace the lookup table that mapped client method names to A2A RPC names with an rpc_method argument passed by each BaseClient method, the same way DefaultRequestHandlerV2 names its methods in @server_plugin_hooks. Interceptors keep receiving the existing method strings; rpc_method is keyword-only and defaults to method, so callers of the private helpers are unaffected.
BasePushNotificationSender now takes a plugins argument (global manager by default) and runs push hooks for every delivery, whichever request handler uses it. This replaces the contextvar that let only DefaultRequestHandlerV2 enable push hooks: not wiring the legacy handler for plugins does not require shared components to block them. ActiveTask is back to its state on main, and a2a.plugins._runtime is removed.
🧪 Code Coverage (vs
|
| Base | PR | Delta | |
|---|---|---|---|
| src/a2a/client/base_client.py | 95.56% | 96.67% | 🟢 +1.11% |
| src/a2a/server/agent_execution/active_task_registry.py | 97.37% | 97.40% | 🟢 +0.03% |
| src/a2a/server/request_handlers/default_request_handler_v2.py | 87.46% | 88.49% | 🟢 +1.03% |
| src/a2a/server/tasks/base_push_notification_sender.py | 95.38% | 94.95% | 🔴 -0.44% |
| src/a2a/server/tasks/task_manager.py | 98.86% | 98.94% | 🟢 +0.08% |
| src/a2a/plugins/base.py (new) | — | 100.00% | — |
| src/a2a/plugins/manager.py (new) | — | 100.00% | — |
| src/a2a/plugins/types.py (new) | — | 100.00% | — |
| src/a2a/server/agent_execution/plugin_agent_executor.py (new) | — | 93.75% | — |
| src/a2a/server/request_handlers/_plugin_hooks.py (new) | — | 94.83% | — |
| Total | 92.96% | 93.19% | 🟢 +0.23% |
Generated by coverage-comment.yml
ServerPlugin's before/after_agent_run run for every request handler method, including ones that never run an agent (task retrieval, push config CRUD), so they become before/after_handler_run. ClientPlugin's before/after_agent_call likewise run for every client request and become before/after_client_request, matching ClientRequest and ClientResponse.
The plugin registry suppresses ruff's PLW0603 on its global statement, and check-spelling flags the rule prefix as an unknown word.
| method='send_message', | ||
| rpc_method='SendMessage', |
There was a problem hiding this comment.
Why do we need both method and rpc_method? Seems like we use different naming for interceptors and plugins. I don't think we need them to have some specific naming (as names are different in different transports), just something that is documented should be enough. Can we somehow reuse method naming format?
There was a problem hiding this comment.
My idea was that client side and server side plugins should get the same method name. But the names of the actual DefaultRequestHandlerV2 methods differ slightly from the client-side "method" values, e.g. on_message_send vs send_message (not message_send). So I added a separate "rpc_method" field and used the same names in the server_plugin_hooks decorators. But yeah that's not great, agree.
I've updated the server side server_plugin_hooks param to match the client "method" values, removed "rpc_method"
Drop the rpc_method argument from BaseClient's call helpers: client plugins now receive the same method name as interceptors, and DefaultRequestHandlerV2 names its methods in @server_plugin_hooks with those same strings, so plugins see one naming scheme on both sides without a second parameter on every call.
Adds plugins to the Python SDK. A plugin is one object that hooks into both the client and
DefaultRequestHandlerV2, so extension points like telemetry, auth or tenant propagation need one registration call instead of per-SDK interceptor wiring.Example:
Interfaces
A plugin implements any of:
SequentialPluginManager runs before-hooks in registration order and after-hooks in reverse; a raising before-hook rejects the call, while after-hook and observer failures are logged and swallowed.
Notes
Tested