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:
committed by
Copybara-Service
parent
780093f389
commit
6f0dcb3e26
@@ -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},
|
||||
}
|
||||
|
||||
@@ -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):
|
||||
|
||||
Reference in New Issue
Block a user