Skip to content

Commit 6f772d2

Browse files
google-genai-botcopybara-github
authored andcommitted
feat: Introduce A2A request interceptors in RemoteA2aAgent
This change adds a new `a2a` subpackage with configuration and utility functions for intercepting requests and responses in `RemoteA2aAgent`. The `RemoteA2aAgent` now accepts an `A2aRemoteAgentConfig` to register `RequestInterceptor` instances, allowing custom logic to be executed before and after the A2A message send. PiperOrigin-RevId: 875559286
1 parent 5f806ed commit 6f772d2

File tree

5 files changed

+405
-6
lines changed

5 files changed

+405
-6
lines changed
Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,25 @@
1+
# Copyright 2026 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
"""A2A agents package."""
16+
17+
from .config import A2aRemoteAgentConfig
18+
from .config import ParametersConfig
19+
from .config import RequestInterceptor
20+
21+
__all__ = [
22+
"A2aRemoteAgentConfig",
23+
"ParametersConfig",
24+
"RequestInterceptor",
25+
]

src/google/adk/a2a/agent/config.py

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,76 @@
1+
# Copyright 2026 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
"""Configuration for A2A agents."""
16+
17+
from __future__ import annotations
18+
19+
from typing import Any
20+
from typing import Awaitable
21+
from typing import Callable
22+
from typing import Optional
23+
from typing import Union
24+
25+
from a2a.client.middleware import ClientCallContext
26+
from a2a.server.events import Event as A2AEvent
27+
from a2a.types import Message as A2AMessage
28+
from a2a.types import MessageSendConfiguration
29+
from pydantic import BaseModel
30+
31+
from ...agents.invocation_context import InvocationContext
32+
from ...events.event import Event
33+
34+
35+
class ParametersConfig(BaseModel):
36+
"""Configuration for the parameters passed to the A2A send_message request."""
37+
38+
request_metadata: Optional[dict[str, Any]] = None
39+
client_call_context: Optional[ClientCallContext] = None
40+
# TODO: Add support for requested_extension and
41+
# message_send_configuration once they are supported by the A2A client.
42+
#
43+
# requested_extension: Optional[list[str]] = None
44+
# message_send_configuration: Optional[MessageSendConfiguration] = None
45+
46+
47+
class RequestInterceptor(BaseModel):
48+
"""Interceptor for A2A requests."""
49+
50+
before_request: Optional[
51+
Callable[
52+
[InvocationContext, A2AMessage, ParametersConfig],
53+
Awaitable[tuple[Union[A2AMessage, Event], ParametersConfig]],
54+
]
55+
] = None
56+
"""Hook executed before the agent starts processing the request.
57+
58+
Returns an Event if the request should be aborted and the Event
59+
returned to the caller.
60+
"""
61+
62+
after_request: Optional[
63+
Callable[
64+
[InvocationContext, A2AEvent, Event], Awaitable[Union[Event, None]]
65+
]
66+
] = None
67+
"""Hook executed after the agent has processed the request.
68+
69+
Returns None if the event should not be sent to the caller.
70+
"""
71+
72+
73+
class A2aRemoteAgentConfig(BaseModel):
74+
"""Configuration for the RemoteA2aAgent."""
75+
76+
request_interceptors: Optional[list[RequestInterceptor]] = None

src/google/adk/a2a/agent/utils.py

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,70 @@
1+
# Copyright 2026 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
"""Utilities for A2A agents."""
16+
17+
from __future__ import annotations
18+
19+
from typing import Optional
20+
from typing import Union
21+
22+
from a2a.client import ClientEvent as A2AClientEvent
23+
from a2a.client.middleware import ClientCallContext
24+
from a2a.types import Message as A2AMessage
25+
26+
from ...agents.invocation_context import InvocationContext
27+
from ...events.event import Event
28+
from .config import ParametersConfig
29+
from .config import RequestInterceptor
30+
31+
32+
async def execute_before_request_interceptors(
33+
request_interceptors: Optional[list[RequestInterceptor]],
34+
ctx: InvocationContext,
35+
a2a_request: A2AMessage,
36+
) -> tuple[Union[A2AMessage, Event], ParametersConfig]:
37+
"""Executes registered before_request interceptors."""
38+
39+
params = ParametersConfig(
40+
client_call_context=ClientCallContext(state=ctx.session.state)
41+
)
42+
if request_interceptors:
43+
for interceptor in request_interceptors:
44+
if not interceptor.before_request:
45+
continue
46+
47+
result, params = await interceptor.before_request(
48+
ctx, a2a_request, params
49+
)
50+
if isinstance(result, Event):
51+
return result, params
52+
a2a_request = result
53+
54+
return a2a_request, params
55+
56+
57+
async def execute_after_request_interceptors(
58+
request_interceptors: Optional[list[RequestInterceptor]],
59+
ctx: InvocationContext,
60+
a2a_response: A2AMessage | A2AClientEvent,
61+
event: Event,
62+
) -> Optional[Event]:
63+
"""Executes registered after_request interceptors."""
64+
if request_interceptors:
65+
for interceptor in reversed(request_interceptors):
66+
if interceptor.after_request:
67+
event = await interceptor.after_request(ctx, a2a_response, event)
68+
if not event:
69+
return None
70+
return event

src/google/adk/agents/remote_a2a_agent.py

Lines changed: 30 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@
3535
from a2a.client.middleware import ClientCallContext
3636
from a2a.types import AgentCard
3737
from a2a.types import Message as A2AMessage
38+
from a2a.types import MessageSendConfiguration
3839
from a2a.types import Part as A2APart
3940
from a2a.types import Role
4041
from a2a.types import TaskArtifactUpdateEvent as A2ATaskArtifactUpdateEvent
@@ -43,13 +44,17 @@
4344
from a2a.types import TransportProtocol as A2ATransport
4445
from google.genai import types as genai_types
4546
import httpx
47+
from pydantic import BaseModel
4648

4749
try:
4850
from a2a.utils.constants import AGENT_CARD_WELL_KNOWN_PATH
4951
except ImportError:
5052
# Fallback for older versions of a2a-sdk.
5153
AGENT_CARD_WELL_KNOWN_PATH = "/.well-known/agent.json"
5254

55+
from ..a2a.agent.config import A2aRemoteAgentConfig
56+
from ..a2a.agent.utils import execute_after_request_interceptors
57+
from ..a2a.agent.utils import execute_before_request_interceptors
5358
from ..a2a.converters.event_converter import convert_a2a_message_to_event
5459
from ..a2a.converters.event_converter import convert_a2a_task_to_event
5560
from ..a2a.converters.event_converter import convert_event_to_a2a_message
@@ -127,6 +132,7 @@ def __init__(
127132
Callable[[InvocationContext, A2AMessage], dict[str, Any]]
128133
] = None,
129134
full_history_when_stateless: bool = False,
135+
config: Optional[A2aRemoteAgentConfig] = None,
130136
**kwargs: Any,
131137
) -> None:
132138
"""Initialize RemoteA2aAgent.
@@ -147,6 +153,7 @@ def __init__(
147153
return Tasks or context IDs) will receive all session events on every
148154
request. If False, the default behavior of sending only events since the
149155
last reply from the agent will be used.
156+
config: Optional configuration object.
150157
**kwargs: Additional arguments passed to BaseAgent
151158
152159
Raises:
@@ -174,6 +181,7 @@ def __init__(
174181
self._a2a_client_factory: Optional[A2AClientFactory] = a2a_client_factory
175182
self._a2a_request_meta_provider = a2a_request_meta_provider
176183
self._full_history_when_stateless = full_history_when_stateless
184+
self._config = config or A2aRemoteAgentConfig()
177185

178186
# Validate and store agent card reference
179187
if isinstance(agent_card, AgentCard):
@@ -558,21 +566,39 @@ async def _run_async_impl(
558566
logger.debug(build_a2a_request_log(a2a_request))
559567

560568
try:
561-
request_metadata = None
569+
a2a_request, parameters = await execute_before_request_interceptors(
570+
self._config.request_interceptors, ctx, a2a_request
571+
)
572+
573+
if isinstance(a2a_request, Event):
574+
yield a2a_request
575+
return
576+
577+
# Backward compatibility
562578
if self._a2a_request_meta_provider:
563-
request_metadata = self._a2a_request_meta_provider(ctx, a2a_request)
579+
parameters.request_metadata = self._a2a_request_meta_provider(
580+
ctx, a2a_request
581+
)
564582

583+
# TODO: Add support for requested_extension and
584+
# message_send_configuration once they are supported by the A2A client.
565585
async for a2a_response in self._a2a_client.send_message(
566586
request=a2a_request,
567-
request_metadata=request_metadata,
568-
context=ClientCallContext(state=ctx.session.state),
587+
request_metadata=parameters.request_metadata,
588+
context=parameters.client_call_context,
569589
):
570590
logger.debug(build_a2a_response_log(a2a_response))
571591

572592
event = await self._handle_a2a_response(a2a_response, ctx)
573593
if not event:
574594
continue
575595

596+
event = await execute_after_request_interceptors(
597+
self._config.request_interceptors, ctx, a2a_response, event
598+
)
599+
if not event:
600+
continue
601+
576602
# Add metadata about the request and response
577603
event.custom_metadata = event.custom_metadata or {}
578604
event.custom_metadata[A2A_METADATA_PREFIX + "request"] = (

0 commit comments

Comments
 (0)