Skip to content

Commit e615a8d

Browse files
authored
feat: add plugin framework (#1352)
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:__ ```python from a2a import plugins class TaskAudit(plugins.TaskObserverPlugin): 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())) ``` __Interfaces__ A plugin implements any of: - ServerPlugin (before_handler_run, after_handler_run) - ClientPlugin (before_client_request, after_client_request) - ExecutionPlugin (before_execution / on_execution_event / after_execution) - TaskObserverPlugin (on_task_event) - PushPlugin (before_push_send, on_push_sent) 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__ - No behavior change when no plugin is registered. - LegacyRequestHandler is not wired for plugins. __Tested__ - Registered plugins that implement every hook: log each call, inject headers.
1 parent 4e133b7 commit e615a8d

18 files changed

Lines changed: 3283 additions & 48 deletions

File tree

‎.github/actions/spelling/allow.txt‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -115,6 +115,7 @@ opensource
115115
otherurl
116116
outerjoin
117117
pb2
118+
PLW
118119
podman
119120
Podman
120121
poolclass

‎src/a2a/client/base_client.py‎

Lines changed: 150 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable
2+
from contextlib import aclosing
23
from typing import Any
34

45
from a2a.client.client import (
@@ -12,6 +13,8 @@
1213
ClientCallInterceptor,
1314
)
1415
from a2a.client.transports.base import ClientTransport
16+
from a2a.plugins.manager import PluginManager, get_global
17+
from a2a.plugins.types import ClientRequest, ClientResponse
1518
from a2a.types.a2a_pb2 import (
1619
AgentCard,
1720
CancelTaskRequest,
@@ -32,20 +35,37 @@
3235

3336

3437
class BaseClient(Client):
35-
"""Base implementation of the A2A client, containing transport-independent logic."""
38+
"""Base implementation of the A2A client, containing transport-independent logic.
39+
40+
Client plugin hooks run around every call: ``before_client_request`` before
41+
the interceptors, ``after_client_request`` after them.
42+
"""
3643

3744
def __init__(
3845
self,
3946
card: AgentCard,
4047
config: ClientConfig,
4148
transport: ClientTransport,
4249
interceptors: list[ClientCallInterceptor],
50+
*,
51+
plugins: PluginManager | None = None,
4352
):
53+
"""Initializes the client.
54+
55+
Args:
56+
card: The card of the agent to call.
57+
config: The client configuration.
58+
transport: The transport used for calls.
59+
interceptors: Interceptors applied to every call.
60+
plugins: Plugin manager whose client hooks run for every call.
61+
Defaults to the global manager at construction time.
62+
"""
4463
super().__init__(interceptors)
4564
self._card = card
4665
self._config = config
4766
self._transport = transport
4867
self._interceptors = interceptors
68+
self._plugins = plugins if plugins is not None else get_global()
4969

5070
async def send_message(
5171
self,
@@ -90,15 +110,20 @@ async def send_message(
90110
yield stream_response
91111
return
92112

93-
async for event in self._execute_stream_with_interceptors(
94-
input_data=request,
95-
method='send_message_streaming',
96-
context=context,
97-
transport_call=lambda req, ctx: (
98-
self._transport.send_message_streaming(req, context=ctx)
99-
),
100-
):
101-
yield event
113+
# aclosing() closes the inner stream (and the transport stream) as
114+
# soon as the caller stops iterating.
115+
async with aclosing(
116+
self._execute_stream_with_interceptors(
117+
input_data=request,
118+
method='send_message_streaming',
119+
context=context,
120+
transport_call=lambda req, ctx: (
121+
self._transport.send_message_streaming(req, context=ctx)
122+
),
123+
)
124+
) as events:
125+
async for event in events:
126+
yield event
102127

103128
def _apply_client_config(self, request: SendMessageRequest) -> None:
104129
request.configuration.return_immediately |= self._config.polling
@@ -324,15 +349,18 @@ async def subscribe(
324349
'client and/or server do not support resubscription.'
325350
)
326351

327-
async for event in self._execute_stream_with_interceptors(
328-
input_data=request,
329-
method='subscribe',
330-
context=context,
331-
transport_call=lambda req, ctx: self._transport.subscribe(
332-
req, context=ctx
333-
),
334-
):
335-
yield event
352+
async with aclosing(
353+
self._execute_stream_with_interceptors(
354+
input_data=request,
355+
method='subscribe',
356+
context=context,
357+
transport_call=lambda req, ctx: self._transport.subscribe(
358+
req, context=ctx
359+
),
360+
)
361+
) as events:
362+
async for event in events:
363+
yield event
336364

337365
async def get_extended_agent_card(
338366
self,
@@ -372,6 +400,19 @@ async def close(self) -> None:
372400
"""Closes the underlying transport."""
373401
await self._transport.close()
374402

403+
def _client_request(
404+
self,
405+
input_data: Any,
406+
method: str,
407+
context: ClientCallContext | None,
408+
) -> ClientRequest:
409+
return ClientRequest(
410+
method=method,
411+
payload=input_data,
412+
agent_card=self._card,
413+
context=context if context is not None else ClientCallContext(),
414+
)
415+
375416
async def _execute_with_interceptors(
376417
self,
377418
input_data: Any,
@@ -380,6 +421,41 @@ async def _execute_with_interceptors(
380421
transport_call: Callable[
381422
[Any, ClientCallContext | None], Awaitable[Any]
382423
],
424+
) -> Any:
425+
"""Runs a unary call through plugins, interceptors and the transport.
426+
427+
``method`` is reported to both plugins and interceptors.
428+
"""
429+
plugins = self._plugins
430+
if not plugins.has_client_hooks:
431+
return await self._call_through_interceptors(
432+
input_data, method, context, transport_call
433+
)
434+
435+
request = self._client_request(input_data, method, context)
436+
await plugins.before_client_request(request)
437+
try:
438+
result = await self._call_through_interceptors(
439+
request.payload, method, request.context, transport_call
440+
)
441+
except Exception as e:
442+
await plugins.after_client_request(
443+
ClientResponse(request.method, None, request.context, error=e)
444+
)
445+
raise
446+
await plugins.after_client_request(
447+
ClientResponse(request.method, result, request.context)
448+
)
449+
return result
450+
451+
async def _call_through_interceptors(
452+
self,
453+
input_data: Any,
454+
method: str,
455+
context: ClientCallContext | None,
456+
transport_call: Callable[
457+
[Any, ClientCallContext | None], Awaitable[Any]
458+
],
383459
) -> Any:
384460
before_args = BeforeArgs(
385461
input=input_data,
@@ -422,8 +498,62 @@ async def _execute_stream_with_interceptors(
422498
transport_call: Callable[
423499
[Any, ClientCallContext | None], AsyncIterator[StreamResponse]
424500
],
425-
) -> AsyncIterator[StreamResponse]:
501+
) -> AsyncGenerator[StreamResponse, None]:
502+
"""Runs a streaming call through plugins, interceptors and the transport.
426503
504+
``method`` is reported to both plugins and interceptors.
505+
"""
506+
plugins = self._plugins
507+
if not plugins.has_client_hooks:
508+
async with aclosing(
509+
self._stream_through_interceptors(
510+
input_data, method, context, transport_call
511+
)
512+
) as events:
513+
async for event in events:
514+
yield event
515+
return
516+
517+
request = self._client_request(input_data, method, context)
518+
await plugins.before_client_request(request)
519+
error: BaseException | None = None
520+
try:
521+
async with aclosing(
522+
self._stream_through_interceptors(
523+
request.payload, method, request.context, transport_call
524+
)
525+
) as events:
526+
async for event in events:
527+
await plugins.after_client_request(
528+
ClientResponse(
529+
request.method,
530+
event,
531+
request.context,
532+
is_stream_event=True,
533+
is_final=False,
534+
)
535+
)
536+
yield event
537+
except Exception as e:
538+
error = e
539+
raise
540+
finally:
541+
# Also runs when the caller stops iterating early.
542+
await plugins.after_client_request(
543+
ClientResponse(
544+
request.method, None, request.context, error=error
545+
)
546+
)
547+
548+
async def _stream_through_interceptors(
549+
self,
550+
input_data: Any,
551+
method: str,
552+
context: ClientCallContext | None,
553+
transport_call: Callable[
554+
[Any, ClientCallContext | None], AsyncIterator[StreamResponse]
555+
],
556+
) -> AsyncGenerator[StreamResponse, None]:
427557
before_args = BeforeArgs(
428558
input=input_data,
429559
method=method,

‎src/a2a/plugins/__init__.py‎

Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,63 @@
1+
"""Plugins: hooks into the A2A client and `DefaultRequestHandlerV2`.
2+
3+
Example:
4+
from a2a import plugins
5+
6+
class TaskAudit(plugins.TaskObserverPlugin):
7+
name = 'task-audit'
8+
9+
async def on_task_event(self, info: plugins.TaskEventInfo) -> None:
10+
print(info.task_id, info.from_state, '->', info.to_state)
11+
12+
plugins.set_global(plugins.sequential(TaskAudit()))
13+
# Create DefaultRequestHandlerV2 and clients after this point.
14+
"""
15+
16+
from a2a.plugins.base import (
17+
ClientPlugin,
18+
CompletePlugin,
19+
ExecutionPlugin,
20+
Plugin,
21+
PushPlugin,
22+
ServerPlugin,
23+
TaskObserverPlugin,
24+
)
25+
from a2a.plugins.manager import (
26+
PluginManager,
27+
SequentialPluginManager,
28+
get_global,
29+
sequential,
30+
set_global,
31+
)
32+
from a2a.plugins.types import (
33+
ClientRequest,
34+
ClientResponse,
35+
PushRequest,
36+
PushResult,
37+
ServerRequest,
38+
ServerResponse,
39+
TaskEventInfo,
40+
)
41+
42+
43+
__all__ = [
44+
'ClientPlugin',
45+
'ClientRequest',
46+
'ClientResponse',
47+
'CompletePlugin',
48+
'ExecutionPlugin',
49+
'Plugin',
50+
'PluginManager',
51+
'PushPlugin',
52+
'PushRequest',
53+
'PushResult',
54+
'SequentialPluginManager',
55+
'ServerPlugin',
56+
'ServerRequest',
57+
'ServerResponse',
58+
'TaskEventInfo',
59+
'TaskObserverPlugin',
60+
'get_global',
61+
'sequential',
62+
'set_global',
63+
]

0 commit comments

Comments
 (0)