feat: Add a new extension for the new version of ADK-A2A integration

This change introduces a new interceptor that adds the 'https://google.github.io/adk-docs/a2a/a2a-extension/' extension to the request headers in the A2A client from the RemoteAgent side. To send this extension along with requests, the RemoteAgent has to be instantiated with the `use_legacy` flag set to False. The AgentExecutor will default to the new implementation when this extension is requested by the client, but this behavior can be disabled via the `use_legacy` flag.
The 'force_new' flag on the agent_executor side can be used to bypass the presence of the extension, and always activate the new version of the agent_executor.

PiperOrigin-RevId: 883021792
This commit is contained in:
Google Team Member
2026-03-13 01:28:52 -07:00
committed by Copybara-Service
parent 780093f389
commit 6f0dcb3e26
7 changed files with 115 additions and 14 deletions
@@ -0,0 +1,13 @@
# Copyright 2026 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
@@ -0,0 +1,56 @@
# Copyright 2026 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Interceptor that injects the new agent version extension."""
from __future__ import annotations
from typing import Union
from a2a.client.middleware import ClientCallContext
from a2a.extensions.common import HTTP_EXTENSION_HEADER
from a2a.types import Message as A2AMessage
from google.adk.a2a.agent.config import ParametersConfig
from google.adk.a2a.agent.config import RequestInterceptor
from google.adk.agents.invocation_context import InvocationContext
from google.adk.events.event import Event
_NEW_A2A_ADK_INTEGRATION_EXTENSION = (
'https://google.github.io/adk-docs/a2a/a2a-extension/'
)
async def _before_request(
_: InvocationContext,
a2a_request: A2AMessage,
params: ParametersConfig,
) -> tuple[Union[A2AMessage, Event], ParametersConfig]:
"""Adds A2A_new_agent_version to client_call_context."""
if params.client_call_context is None:
params.client_call_context = ClientCallContext()
http_kwargs = params.client_call_context.state.get('http_kwargs', {})
headers = http_kwargs.get('headers', {})
a2a_extensions = headers.get(HTTP_EXTENSION_HEADER, '').split(',')
a2a_extensions = [ext for ext in a2a_extensions if ext]
if _NEW_A2A_ADK_INTEGRATION_EXTENSION not in a2a_extensions:
a2a_extensions.append(_NEW_A2A_ADK_INTEGRATION_EXTENSION)
headers[HTTP_EXTENSION_HEADER] = ','.join(a2a_extensions)
http_kwargs['headers'] = headers
params.client_call_context.state['http_kwargs'] = http_kwargs
return a2a_request, params
_new_integration_extension_interceptor = RequestInterceptor(
before_request=_before_request
)
@@ -39,6 +39,7 @@ from google.adk.runners import Runner
from typing_extensions import override
from ...utils.context_utils import Aclosing
from ..agent.interceptors.new_integration_extension import _NEW_A2A_ADK_INTEGRATION_EXTENSION
from ..converters.request_converter import AgentRunRequest
from ..converters.utils import _get_adk_metadata_key
from ..experimental import a2a_experimental
@@ -62,7 +63,9 @@ class A2aAgentExecutor(AgentExecutor):
Args:
runner: The runner to use for the agent.
config: The config to use for the executor.
use_legacy: Whether to use the legacy executor implementation.
use_legacy: If true, force the legacy implementation.
force_new_version: If true, force the new implementation regardless of the
extension.
"""
def __init__(
@@ -70,15 +73,15 @@ class A2aAgentExecutor(AgentExecutor):
*,
runner: Runner | Callable[..., Runner | Awaitable[Runner]],
config: Optional[A2aAgentExecutorConfig] = None,
use_legacy: bool = True,
use_legacy: bool = False,
force_new_version: bool = False,
):
super().__init__()
if not use_legacy:
self._executor_impl = ExecutorImpl(runner=runner, config=config)
else:
self._executor_impl = None
self._runner = runner
self._config = config or A2aAgentExecutorConfig()
self._runner = runner
self._config = config or A2aAgentExecutorConfig()
self._use_legacy = use_legacy
self._force_new_version = force_new_version
self._executor_impl = None
async def _resolve_runner(self) -> Runner:
"""Resolve the runner, handling cases where it's a callable that returns a Runner."""
@@ -129,7 +132,16 @@ class A2aAgentExecutor(AgentExecutor):
* Converts the ADK output events into A2A task updates
* Publishes the updates back to A2A server via event queue
"""
if self._executor_impl:
should_use_new_impl = not self._use_legacy and (
self._force_new_version or self._check_new_version_extension(context)
)
if should_use_new_impl:
if self._executor_impl is None:
self._executor_impl = ExecutorImpl(
runner=self._runner,
config=self._config,
)
await self._executor_impl.execute(context, event_queue)
return
@@ -338,3 +350,10 @@ class A2aAgentExecutor(AgentExecutor):
run_request.session_id = session.id
return session
def _check_new_version_extension(self, context: RequestContext):
"""Check if the extension for the new version is requested and activate it."""
if _NEW_A2A_ADK_INTEGRATION_EXTENSION in context.requested_extensions:
context.add_activated_extension(_NEW_A2A_ADK_INTEGRATION_EXTENSION)
return True
return False
@@ -39,6 +39,7 @@ from typing_extensions import override
from ...runners import Runner
from ...utils.context_utils import Aclosing
from ..agent.interceptors.new_integration_extension import _NEW_A2A_ADK_INTEGRATION_EXTENSION
from ..converters.from_adk_event import create_error_status_event
from ..converters.long_running_functions import handle_user_input
from ..converters.long_running_functions import LongRunningFunctions
@@ -306,5 +307,5 @@ class _A2aAgentExecutor(AgentExecutor):
_get_adk_metadata_key('session_id'): executor_context.session_id,
# TODO: Remove this metadata once the new agent executor
# is fully adopted.
_get_adk_metadata_key('agent_executor_v2'): True,
_NEW_A2A_ADK_INTEGRATION_EXTENSION: {'adk_agent_executor_v2': True},
}
+13 -3
View File
@@ -54,6 +54,8 @@ except ImportError:
AGENT_CARD_WELL_KNOWN_PATH = "/.well-known/agent.json"
from ..a2a.agent.config import A2aRemoteAgentConfig
from ..a2a.agent.interceptors.new_integration_extension import _NEW_A2A_ADK_INTEGRATION_EXTENSION
from ..a2a.agent.interceptors.new_integration_extension import _new_integration_extension_interceptor
from ..a2a.agent.utils import execute_after_request_interceptors
from ..a2a.agent.utils import execute_before_request_interceptors
from ..a2a.converters.event_converter import convert_a2a_message_to_event
@@ -135,6 +137,7 @@ class RemoteA2aAgent(BaseAgent):
] = None,
full_history_when_stateless: bool = False,
config: Optional[A2aRemoteAgentConfig] = None,
use_legacy: bool = True,
**kwargs: Any,
) -> None:
"""Initialize RemoteA2aAgent.
@@ -156,6 +159,8 @@ class RemoteA2aAgent(BaseAgent):
request. If False, the default behavior of sending only events since the
last reply from the agent will be used.
config: Optional configuration object.
use_legacy: If false, send request to the server including the extension
indicating that the server should use the new implementation.
**kwargs: Additional arguments passed to BaseAgent
Raises:
@@ -185,6 +190,13 @@ class RemoteA2aAgent(BaseAgent):
self._full_history_when_stateless = full_history_when_stateless
self._config = config or A2aRemoteAgentConfig()
if not use_legacy:
if self._config.request_interceptors is None:
self._config.request_interceptors = []
self._config.request_interceptors.append(
_new_integration_extension_interceptor
)
# Validate and store agent card reference
if isinstance(agent_card, AgentCard):
self._agent_card = agent_card
@@ -669,9 +681,7 @@ class RemoteA2aAgent(BaseAgent):
else:
metadata = a2a_response.metadata
if metadata and metadata.get(
_get_adk_metadata_key("agent_executor_v2")
):
if metadata and metadata.get(_NEW_A2A_ADK_INTEGRATION_EXTENSION):
event = await self._handle_a2a_response_v2(a2a_response, ctx)
else:
event = await self._handle_a2a_response(a2a_response, ctx)
@@ -66,6 +66,7 @@ class TestA2aAgentExecutor:
self.mock_context.current_task = None
self.mock_context.task_id = "test-task-id"
self.mock_context.context_id = "test-context-id"
self.mock_context.requested_extensions = []
self.mock_event_queue = Mock(spec=EventQueue)
@@ -29,6 +29,7 @@ from a2a.types import TextPart
from google.adk.a2a.converters.request_converter import AgentRunRequest
from google.adk.a2a.converters.utils import _get_adk_metadata_key
from google.adk.a2a.executor.a2a_agent_executor_impl import _A2aAgentExecutor as A2aAgentExecutor
from google.adk.a2a.executor.a2a_agent_executor_impl import _NEW_A2A_ADK_INTEGRATION_EXTENSION
from google.adk.a2a.executor.a2a_agent_executor_impl import A2aAgentExecutorConfig
from google.adk.a2a.executor.config import ExecuteInterceptor
from google.adk.events.event import Event
@@ -77,7 +78,7 @@ class TestA2aAgentExecutor:
_get_adk_metadata_key("app_name"): "test-app",
_get_adk_metadata_key("user_id"): "test-user",
_get_adk_metadata_key("session_id"): "test-session",
_get_adk_metadata_key("agent_executor_v2"): True,
_NEW_A2A_ADK_INTEGRATION_EXTENSION: {"adk_agent_executor_v2": True},
}
async def _create_async_generator(self, items):