From ebd3f610f7e7e6dcd44137e1c3799403e61d399a Mon Sep 17 00:00:00 2001 From: droideronline Date: Thu, 8 Oct 2026 14:04:41 +0530 Subject: [PATCH 1/5] Python: A2AExecutor: add prepare_session hook Add an overridable hook called after the session is created and before the agent runs so applications can initialize session state from the inbound request. --- .../a2a/agent_framework_a2a/_a2a_executor.py | 16 +++ .../packages/a2a/tests/test_a2a_executor.py | 109 ++++++++++++++++++ 2 files changed, 125 insertions(+) diff --git a/python/packages/a2a/agent_framework_a2a/_a2a_executor.py b/python/packages/a2a/agent_framework_a2a/_a2a_executor.py index 2aa028d3856..13df39d1508 100644 --- a/python/packages/a2a/agent_framework_a2a/_a2a_executor.py +++ b/python/packages/a2a/agent_framework_a2a/_a2a_executor.py @@ -164,6 +164,7 @@ async def execute(self, context: RequestContext, event_queue: EventQueue) -> Non await updater.start_work() session = self._agent.create_session(session_id=task.context_id) + await self.prepare_session(context, session) if self._stream: await self._run_stream(query, session, updater) @@ -181,6 +182,21 @@ async def execute(self, context: RequestContext, event_queue: EventQueue) -> Non message=updater.new_agent_message([Part(text=str(e))]), ) + async def prepare_session(self, context: RequestContext, session: AgentSession) -> None: + """Initialize the agent session from the inbound A2A request. + + Called after the session is created and before the agent runs. The default + implementation does nothing. Override it to copy trusted values from the request, + for example ``context.metadata`` (``SendMessageRequest.metadata``), into + ``session.state`` so middleware and tools on the hosted agent can read them. + Nothing from the request is applied automatically, so the application decides + which keys to trust. + + Args: + context: The inbound A2A request context. + session: The session the agent will run with. + """ + async def _run_stream(self, query: Any, session: AgentSession, updater: TaskUpdater) -> None: """Run the agent in streaming mode and publish updates to the task updater.""" response_stream = self._agent.run(query, session=session, stream=True, **self._run_kwargs) diff --git a/python/packages/a2a/tests/test_a2a_executor.py b/python/packages/a2a/tests/test_a2a_executor.py index 062cb40b851..a0859e27f96 100644 --- a/python/packages/a2a/tests/test_a2a_executor.py +++ b/python/packages/a2a/tests/test_a2a_executor.py @@ -7,6 +7,7 @@ from a2a.types import Part, Task, TaskState from agent_framework import ( AgentResponseUpdate, + AgentSession, Content, Message, SupportsAgentRun, @@ -923,3 +924,111 @@ async def test_full_execution_flow_with_responses( mock_updater.start_work.assert_called_once() cast(Any, executor.handle_events).assert_called_once() mock_updater.complete.assert_called_once() + + +class TestA2AExecutorPrepareSession: + """Tests for the prepare_session hook.""" + + async def _execute( + self, + executor: A2AExecutor, + mock_request_context: MagicMock, + mock_event_queue: MagicMock, + mock_task: Task, + session: Any, + run: Any = None, + ) -> MagicMock: + mock_request_context.get_user_input = MagicMock(return_value="Hello") + mock_request_context.current_task = mock_task + mock_request_context.context_id = "ctx-123" + mock_request_context.message = MagicMock() + mock_request_context.metadata = {"user_id": "42", "user_name": "Dinesh"} + + response = MagicMock(spec=AgentResponse) + response.messages = [Message(role="assistant", contents=[Content.from_text(text="Hi")])] + cast(Any, executor._agent).run = AsyncMock(side_effect=run) if run else AsyncMock(return_value=response) + cast(Any, executor._agent).create_session = MagicMock(return_value=session) + + with patch("agent_framework_a2a._a2a_executor.TaskUpdater") as mock_updater_class: + mock_updater = MagicMock() + mock_updater.submit = AsyncMock() + mock_updater.start_work = AsyncMock() + mock_updater.complete = AsyncMock() + mock_updater.update_status = AsyncMock() + mock_updater_class.return_value = mock_updater + await executor.execute(mock_request_context, mock_event_queue) + return mock_updater + + async def test_default_prepare_session_does_not_touch_session( + self, + executor: A2AExecutor, + mock_request_context: MagicMock, + mock_event_queue: MagicMock, + mock_task: Task, + ) -> None: + session = AgentSession() + + await self._execute(executor, mock_request_context, mock_event_queue, mock_task, session) + + assert session.state == {} + assert cast(Any, executor._agent.run).call_args.kwargs["session"] is session + + async def test_overridden_prepare_session_runs_before_agent( + self, + mock_agent: MagicMock, + mock_request_context: MagicMock, + mock_event_queue: MagicMock, + mock_task: Task, + ) -> None: + seen_in_run: dict[str, Any] = {} + + class _Executor(A2AExecutor): + async def prepare_session(self, context: Any, session: AgentSession) -> None: + session.state["user_id"] = context.metadata.get("user_id") + session.state["user_name"] = context.metadata.get("user_name") + + async def _run(*args: Any, session: AgentSession, **kwargs: Any) -> Any: + seen_in_run.update(session.state) + response = MagicMock(spec=AgentResponse) + response.messages = [Message(role="assistant", contents=[Content.from_text(text="Hi")])] + return response + + executor = _Executor(mock_agent) + session = AgentSession() + mock_updater = await self._execute( + executor, mock_request_context, mock_event_queue, mock_task, session, run=_run + ) + assert session.state == {"user_id": "42", "user_name": "Dinesh"} + assert seen_in_run == {"user_id": "42", "user_name": "Dinesh"} + mock_updater.complete.assert_called_once() + + async def test_prepare_session_error_fails_the_task( + self, + mock_agent: MagicMock, + mock_request_context: MagicMock, + mock_event_queue: MagicMock, + mock_task: Task, + ) -> None: + class _Executor(A2AExecutor): + async def prepare_session(self, context: Any, session: AgentSession) -> None: + raise ValueError("untrusted metadata") + + executor = _Executor(mock_agent) + mock_updater = MagicMock() + mock_updater.submit = AsyncMock() + mock_updater.start_work = AsyncMock() + mock_updater.complete = AsyncMock() + mock_updater.update_status = AsyncMock() + mock_updater.new_agent_message = MagicMock(return_value="message_obj") + mock_request_context.get_user_input = MagicMock(return_value="Hello") + mock_request_context.current_task = mock_task + mock_request_context.context_id = "ctx-123" + mock_request_context.message = MagicMock() + cast(Any, mock_agent).create_session = MagicMock(return_value=AgentSession()) + + with patch("agent_framework_a2a._a2a_executor.TaskUpdater", return_value=mock_updater): + await executor.execute(mock_request_context, mock_event_queue) + + mock_agent.run.assert_not_called() + mock_updater.complete.assert_not_called() + assert mock_updater.update_status.call_args.kwargs["state"] == TaskState.TASK_STATE_FAILED From 98d6af70ca42c815e985dcbf0767ebce33e87c23 Mon Sep 17 00:00:00 2001 From: droideronline Date: Thu, 8 Oct 2026 14:18:41 +0530 Subject: [PATCH 2/5] Python: A2AExecutor: accept forwarded run kwargs from request metadata Replace the prepare_session hook with opt-in acceptance of selected function_invocation_kwargs and client_kwargs keys read from the agent_framework key of SendMessageRequest.metadata. --- .../a2a/agent_framework_a2a/_a2a_executor.py | 84 +++++--- .../packages/a2a/tests/test_a2a_executor.py | 191 ++++++++++-------- 2 files changed, 161 insertions(+), 114 deletions(-) diff --git a/python/packages/a2a/agent_framework_a2a/_a2a_executor.py b/python/packages/a2a/agent_framework_a2a/_a2a_executor.py index 13df39d1508..763e3b360aa 100644 --- a/python/packages/a2a/agent_framework_a2a/_a2a_executor.py +++ b/python/packages/a2a/agent_framework_a2a/_a2a_executor.py @@ -4,9 +4,9 @@ import logging import uuid from asyncio import CancelledError -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from functools import partial -from typing import Any +from typing import Any, cast from a2a.helpers import new_task_from_user_message from a2a.server.agent_execution import AgentExecutor, RequestContext @@ -27,6 +27,12 @@ logger = logging.getLogger("agent_framework.a2a") +# Key in ``SendMessageRequest.metadata`` that carries forwarded ``function_invocation_kwargs`` and ``client_kwargs``. +AGENT_FRAMEWORK_METADATA_KEY = "agent_framework" +_FORWARDED_KWARGS_NAMES = ("function_invocation_kwargs", "client_kwargs") +# Framework-managed keys that a remote caller can never set. +_RESERVED_KWARGS = frozenset({"session", "middleware"}) + class A2AExecutor(AgentExecutor): """Execute AI agents using the A2A (Agent-to-Agent) protocol. @@ -89,9 +95,19 @@ class A2AExecutor(AgentExecutor): agent: The AI agent to execute. stream: Whether to stream the agent response. Defaults to False. run_kwargs: Additional keyword arguments to pass to the agent's run method. + accepted_kwargs: Names of keys the executor accepts from the caller's request. A client that + forwards ``function_invocation_kwargs`` and ``client_kwargs`` (see ``A2AAgent``) sends them + in the request metadata, and only the keys named here are passed to the agent's run + method. The default (``None``) accepts nothing. """ - def __init__(self, agent: SupportsAgentRun, stream: bool = False, run_kwargs: Mapping[str, Any] | None = None): + def __init__( + self, + agent: SupportsAgentRun, + stream: bool = False, + run_kwargs: Mapping[str, Any] | None = None, + accepted_kwargs: Sequence[str] | None = None, + ): """Initialize the A2AExecutor with the specified agent. Args: @@ -99,6 +115,10 @@ def __init__(self, agent: SupportsAgentRun, stream: bool = False, run_kwargs: Ma stream: Whether to stream the agent response. Defaults to False. run_kwargs: Additional keyword arguments to pass to the agent's run method. Cannot contain 'session' or 'stream' as these are managed by the executor. + accepted_kwargs: Names of keys to accept from the request's forwarded + ``function_invocation_kwargs`` and ``client_kwargs``. Accepted values are passed to the + agent's run method. Values configured in ``run_kwargs`` take precedence on conflict. + Defaults to accepting nothing. Raises: ValueError: If run_kwargs contains 'session' or 'stream'. @@ -112,6 +132,7 @@ def __init__(self, agent: SupportsAgentRun, stream: bool = False, run_kwargs: Ma if "stream" in run_kwargs: raise ValueError("run_kwargs cannot contain 'stream' as it is managed by the executor.") self._run_kwargs: Mapping[str, Any] = run_kwargs or {} + self._accepted_kwargs: frozenset[str] = frozenset(accepted_kwargs or ()) - _RESERVED_KWARGS @override async def cancel(self, context: RequestContext, event_queue: EventQueue) -> None: @@ -164,12 +185,12 @@ async def execute(self, context: RequestContext, event_queue: EventQueue) -> Non await updater.start_work() session = self._agent.create_session(session_id=task.context_id) - await self.prepare_session(context, session) + run_kwargs = self._build_run_kwargs(context) if self._stream: - await self._run_stream(query, session, updater) + await self._run_stream(query, session, updater, run_kwargs) else: - await self._run(query, session, updater) + await self._run(query, session, updater, run_kwargs) # Mark as complete await updater.complete() @@ -182,24 +203,33 @@ async def execute(self, context: RequestContext, event_queue: EventQueue) -> Non message=updater.new_agent_message([Part(text=str(e))]), ) - async def prepare_session(self, context: RequestContext, session: AgentSession) -> None: - """Initialize the agent session from the inbound A2A request. - - Called after the session is created and before the agent runs. The default - implementation does nothing. Override it to copy trusted values from the request, - for example ``context.metadata`` (``SendMessageRequest.metadata``), into - ``session.state`` so middleware and tools on the hosted agent can read them. - Nothing from the request is applied automatically, so the application decides - which keys to trust. - - Args: - context: The inbound A2A request context. - session: The session the agent will run with. - """ - - async def _run_stream(self, query: Any, session: AgentSession, updater: TaskUpdater) -> None: + def _build_run_kwargs(self, context: RequestContext) -> dict[str, Any]: + """Merge the accepted request kwargs under the configured ``run_kwargs``.""" + run_kwargs: dict[str, Any] = dict(self._run_kwargs) + if not self._accepted_kwargs: + return run_kwargs + forwarded = context.metadata.get(AGENT_FRAMEWORK_METADATA_KEY) + if not isinstance(forwarded, Mapping): + return run_kwargs + for name in _FORWARDED_KWARGS_NAMES: + values = cast("Mapping[str, Any]", forwarded).get(name) + if not isinstance(values, Mapping): + continue + accepted = { + key: value for key, value in cast("Mapping[str, Any]", values).items() if key in self._accepted_kwargs + } + if accepted: + configured = run_kwargs.get(name) + run_kwargs[name] = {**accepted, **(configured or {})} + return run_kwargs + + async def _run_stream( + self, query: Any, session: AgentSession, updater: TaskUpdater, run_kwargs: Mapping[str, Any] | None = None + ) -> None: """Run the agent in streaming mode and publish updates to the task updater.""" - response_stream = self._agent.run(query, session=session, stream=True, **self._run_kwargs) + response_stream = self._agent.run( + query, session=session, stream=True, **(self._run_kwargs if run_kwargs is None else run_kwargs) + ) streamed_artifact_ids: set[str] = set() # Generate a stable artifact ID for the entire stream so all chunks share the same ID. # This ensures clients can coalesce streaming tokens into a single artifact/message @@ -216,9 +246,13 @@ async def _run_stream(self, query: Any, session: AgentSession, updater: TaskUpda ) ).get_final_response() - async def _run(self, query: Any, session: AgentSession, updater: TaskUpdater) -> None: + async def _run( + self, query: Any, session: AgentSession, updater: TaskUpdater, run_kwargs: Mapping[str, Any] | None = None + ) -> None: """Run the agent in non-streaming mode and publish messages to the task updater.""" - response = await self._agent.run(query, session=session, stream=False, **self._run_kwargs) + response = await self._agent.run( + query, session=session, stream=False, **(self._run_kwargs if run_kwargs is None else run_kwargs) + ) response_messages = response.messages if not isinstance(response_messages, list): diff --git a/python/packages/a2a/tests/test_a2a_executor.py b/python/packages/a2a/tests/test_a2a_executor.py index a0859e27f96..0c9f9d59d61 100644 --- a/python/packages/a2a/tests/test_a2a_executor.py +++ b/python/packages/a2a/tests/test_a2a_executor.py @@ -7,14 +7,13 @@ from a2a.types import Part, Task, TaskState from agent_framework import ( AgentResponseUpdate, - AgentSession, Content, Message, SupportsAgentRun, ) from agent_framework._types import AgentResponse from agent_framework.a2a import A2AExecutor -from pytest import fixture, raises +from pytest import fixture, mark, raises @fixture @@ -926,109 +925,123 @@ async def test_full_execution_flow_with_responses( mock_updater.complete.assert_called_once() -class TestA2AExecutorPrepareSession: - """Tests for the prepare_session hook.""" +class TestA2AExecutorForwardedKwargs: + """Tests for accepting forwarded function_invocation_kwargs and client_kwargs.""" - async def _execute( - self, - executor: A2AExecutor, - mock_request_context: MagicMock, - mock_event_queue: MagicMock, - mock_task: Task, - session: Any, - run: Any = None, - ) -> MagicMock: - mock_request_context.get_user_input = MagicMock(return_value="Hello") - mock_request_context.current_task = mock_task - mock_request_context.context_id = "ctx-123" - mock_request_context.message = MagicMock() - mock_request_context.metadata = {"user_id": "42", "user_name": "Dinesh"} + @staticmethod + async def _execute(executor: A2AExecutor, metadata: Any, mock_event_queue: MagicMock, mock_task: Task) -> MagicMock: + request_context = MagicMock() + request_context.get_user_input = MagicMock(return_value="Hello") + request_context.current_task = mock_task + request_context.context_id = "ctx-123" + request_context.message = MagicMock() + request_context.metadata = metadata response = MagicMock(spec=AgentResponse) response.messages = [Message(role="assistant", contents=[Content.from_text(text="Hi")])] - cast(Any, executor._agent).run = AsyncMock(side_effect=run) if run else AsyncMock(return_value=response) - cast(Any, executor._agent).create_session = MagicMock(return_value=session) + cast(Any, executor._agent).run = AsyncMock(return_value=response) + cast(Any, executor._agent).create_session = MagicMock() with patch("agent_framework_a2a._a2a_executor.TaskUpdater") as mock_updater_class: mock_updater = MagicMock() - mock_updater.submit = AsyncMock() - mock_updater.start_work = AsyncMock() - mock_updater.complete = AsyncMock() - mock_updater.update_status = AsyncMock() + for name in ("submit", "start_work", "complete", "update_status"): + setattr(mock_updater, name, AsyncMock()) mock_updater_class.return_value = mock_updater - await executor.execute(mock_request_context, mock_event_queue) - return mock_updater + await executor.execute(request_context, mock_event_queue) + return cast(MagicMock, executor._agent.run) - async def test_default_prepare_session_does_not_touch_session( - self, - executor: A2AExecutor, - mock_request_context: MagicMock, - mock_event_queue: MagicMock, - mock_task: Task, + @staticmethod + def _forwarded(**kwargs: Any) -> dict[str, Any]: + return {"agent_framework": kwargs} + + async def test_accepted_keys_are_passed_to_agent_run( + self, mock_agent: MagicMock, mock_event_queue: MagicMock, mock_task: Task ) -> None: - session = AgentSession() + executor = A2AExecutor(mock_agent, accepted_kwargs=["tenant", "trace"]) + metadata = self._forwarded( + function_invocation_kwargs={"tenant": "acme", "other": 1}, + client_kwargs={"trace": {"id": "abc"}, "other": 2}, + ) - await self._execute(executor, mock_request_context, mock_event_queue, mock_task, session) + run = await self._execute(executor, metadata, mock_event_queue, mock_task) - assert session.state == {} - assert cast(Any, executor._agent.run).call_args.kwargs["session"] is session + kwargs = run.call_args.kwargs + assert kwargs["function_invocation_kwargs"] == {"tenant": "acme"} + assert kwargs["client_kwargs"] == {"trace": {"id": "abc"}} - async def test_overridden_prepare_session_runs_before_agent( - self, - mock_agent: MagicMock, - mock_request_context: MagicMock, - mock_event_queue: MagicMock, - mock_task: Task, + async def test_nothing_is_accepted_by_default( + self, mock_agent: MagicMock, mock_event_queue: MagicMock, mock_task: Task ) -> None: - seen_in_run: dict[str, Any] = {} - - class _Executor(A2AExecutor): - async def prepare_session(self, context: Any, session: AgentSession) -> None: - session.state["user_id"] = context.metadata.get("user_id") - session.state["user_name"] = context.metadata.get("user_name") - - async def _run(*args: Any, session: AgentSession, **kwargs: Any) -> Any: - seen_in_run.update(session.state) - response = MagicMock(spec=AgentResponse) - response.messages = [Message(role="assistant", contents=[Content.from_text(text="Hi")])] - return response - - executor = _Executor(mock_agent) - session = AgentSession() - mock_updater = await self._execute( - executor, mock_request_context, mock_event_queue, mock_task, session, run=_run + executor = A2AExecutor(mock_agent) + metadata = self._forwarded(function_invocation_kwargs={"tenant": "acme"}, client_kwargs={"trace": "abc"}) + + run = await self._execute(executor, metadata, mock_event_queue, mock_task) + + assert "function_invocation_kwargs" not in run.call_args.kwargs + assert "client_kwargs" not in run.call_args.kwargs + + async def test_configured_run_kwargs_win_on_conflict( + self, mock_agent: MagicMock, mock_event_queue: MagicMock, mock_task: Task + ) -> None: + executor = A2AExecutor( + mock_agent, + run_kwargs={"function_invocation_kwargs": {"tenant": "server", "fixed": True}}, + accepted_kwargs=["tenant", "extra"], ) - assert session.state == {"user_id": "42", "user_name": "Dinesh"} - assert seen_in_run == {"user_id": "42", "user_name": "Dinesh"} - mock_updater.complete.assert_called_once() + metadata = self._forwarded(function_invocation_kwargs={"tenant": "caller", "extra": "x"}) - async def test_prepare_session_error_fails_the_task( - self, - mock_agent: MagicMock, - mock_request_context: MagicMock, - mock_event_queue: MagicMock, - mock_task: Task, + run = await self._execute(executor, metadata, mock_event_queue, mock_task) + + assert run.call_args.kwargs["function_invocation_kwargs"] == {"tenant": "server", "fixed": True, "extra": "x"} + + async def test_reserved_keys_are_never_accepted( + self, mock_agent: MagicMock, mock_event_queue: MagicMock, mock_task: Task ) -> None: - class _Executor(A2AExecutor): - async def prepare_session(self, context: Any, session: AgentSession) -> None: - raise ValueError("untrusted metadata") - - executor = _Executor(mock_agent) - mock_updater = MagicMock() - mock_updater.submit = AsyncMock() - mock_updater.start_work = AsyncMock() - mock_updater.complete = AsyncMock() - mock_updater.update_status = AsyncMock() - mock_updater.new_agent_message = MagicMock(return_value="message_obj") - mock_request_context.get_user_input = MagicMock(return_value="Hello") - mock_request_context.current_task = mock_task - mock_request_context.context_id = "ctx-123" - mock_request_context.message = MagicMock() - cast(Any, mock_agent).create_session = MagicMock(return_value=AgentSession()) + executor = A2AExecutor(mock_agent, accepted_kwargs=["session", "middleware", "tenant"]) + metadata = self._forwarded( + client_kwargs={"session": "x", "middleware": ["y"], "tenant": "acme"}, + ) - with patch("agent_framework_a2a._a2a_executor.TaskUpdater", return_value=mock_updater): - await executor.execute(mock_request_context, mock_event_queue) + run = await self._execute(executor, metadata, mock_event_queue, mock_task) + + assert run.call_args.kwargs["client_kwargs"] == {"tenant": "acme"} + + @mark.parametrize( + "metadata", + [{}, {"agent_framework": "text"}, {"agent_framework": {"client_kwargs": "text"}}], + ids=["missing", "not-a-mapping", "values-not-a-mapping"], + ) + async def test_malformed_metadata_is_ignored( + self, mock_agent: MagicMock, mock_event_queue: MagicMock, mock_task: Task, metadata: Any + ) -> None: + executor = A2AExecutor(mock_agent, accepted_kwargs=["tenant"]) + + run = await self._execute(executor, metadata, mock_event_queue, mock_task) + + assert "client_kwargs" not in run.call_args.kwargs + assert "function_invocation_kwargs" not in run.call_args.kwargs + + async def test_accepted_kwargs_are_passed_when_streaming( + self, mock_agent: MagicMock, mock_event_queue: MagicMock, mock_task: Task + ) -> None: + executor = A2AExecutor(mock_agent, stream=True, accepted_kwargs=["tenant"]) + stream = MagicMock() + stream.with_transform_hook = MagicMock(return_value=stream) + stream.get_final_response = AsyncMock() + request_context = MagicMock() + request_context.get_user_input = MagicMock(return_value="Hello") + request_context.current_task = mock_task + request_context.context_id = "ctx-123" + request_context.message = MagicMock() + request_context.metadata = self._forwarded(function_invocation_kwargs={"tenant": "acme"}) + cast(Any, mock_agent).run = MagicMock(return_value=stream) + cast(Any, mock_agent).create_session = MagicMock() + + with patch("agent_framework_a2a._a2a_executor.TaskUpdater") as mock_updater_class: + mock_updater = MagicMock() + for name in ("submit", "start_work", "complete", "update_status"): + setattr(mock_updater, name, AsyncMock()) + mock_updater_class.return_value = mock_updater + await executor.execute(request_context, mock_event_queue) - mock_agent.run.assert_not_called() - mock_updater.complete.assert_not_called() - assert mock_updater.update_status.call_args.kwargs["state"] == TaskState.TASK_STATE_FAILED + assert mock_agent.run.call_args.kwargs["function_invocation_kwargs"] == {"tenant": "acme"} From 235184cb691b84b4ab50c9a3d59d52bc6a94991b Mon Sep 17 00:00:00 2001 From: droideronline Date: Fri, 9 Oct 2026 22:02:08 +0530 Subject: [PATCH 3/5] Python: A2AExecutor: reject framework-owned names in accepted_kwargs Fail fast at construction when accepted_kwargs contains a name the framework owns (stream, messages, options, session, middleware and similar) or an underscore-prefixed name, so a remote caller cannot cause duplicate keyword errors. --- .../a2a/agent_framework_a2a/_a2a_executor.py | 27 +++++++++++++++---- .../packages/a2a/tests/test_a2a_executor.py | 18 +++++-------- 2 files changed, 29 insertions(+), 16 deletions(-) diff --git a/python/packages/a2a/agent_framework_a2a/_a2a_executor.py b/python/packages/a2a/agent_framework_a2a/_a2a_executor.py index 763e3b360aa..c0497828112 100644 --- a/python/packages/a2a/agent_framework_a2a/_a2a_executor.py +++ b/python/packages/a2a/agent_framework_a2a/_a2a_executor.py @@ -30,8 +30,19 @@ # Key in ``SendMessageRequest.metadata`` that carries forwarded ``function_invocation_kwargs`` and ``client_kwargs``. AGENT_FRAMEWORK_METADATA_KEY = "agent_framework" _FORWARDED_KWARGS_NAMES = ("function_invocation_kwargs", "client_kwargs") -# Framework-managed keys that a remote caller can never set. -_RESERVED_KWARGS = frozenset({"session", "middleware"}) +# Names owned by the framework (chat client and tool invocation parameters). A remote caller must never be able to +# set them, since a duplicate keyword would fail the task or override framework behavior. +_RESERVED_KWARGS = frozenset({ + "messages", + "stream", + "options", + "session", + "middleware", + "compaction_strategy", + "tokenizer", + "function_invocation_kwargs", + "client_kwargs", +}) class A2AExecutor(AgentExecutor): @@ -118,10 +129,13 @@ def __init__( accepted_kwargs: Names of keys to accept from the request's forwarded ``function_invocation_kwargs`` and ``client_kwargs``. Accepted values are passed to the agent's run method. Values configured in ``run_kwargs`` take precedence on conflict. - Defaults to accepting nothing. + Defaults to accepting nothing. Framework-owned names (for example ``stream``, + ``messages``, ``options`` and ``session``) and names starting with an underscore are + rejected. Raises: - ValueError: If run_kwargs contains 'session' or 'stream'. + ValueError: If run_kwargs contains 'session' or 'stream', or accepted_kwargs contains a + framework-owned name. """ super().__init__() self._agent: SupportsAgentRun = agent @@ -132,7 +146,10 @@ def __init__( if "stream" in run_kwargs: raise ValueError("run_kwargs cannot contain 'stream' as it is managed by the executor.") self._run_kwargs: Mapping[str, Any] = run_kwargs or {} - self._accepted_kwargs: frozenset[str] = frozenset(accepted_kwargs or ()) - _RESERVED_KWARGS + reserved = sorted(key for key in accepted_kwargs or () if key in _RESERVED_KWARGS or key.startswith("_")) + if reserved: + raise ValueError(f"accepted_kwargs cannot contain framework-owned names: {', '.join(reserved)}.") + self._accepted_kwargs: frozenset[str] = frozenset(accepted_kwargs or ()) @override async def cancel(self, context: RequestContext, event_queue: EventQueue) -> None: diff --git a/python/packages/a2a/tests/test_a2a_executor.py b/python/packages/a2a/tests/test_a2a_executor.py index 0c9f9d59d61..fce37fe84de 100644 --- a/python/packages/a2a/tests/test_a2a_executor.py +++ b/python/packages/a2a/tests/test_a2a_executor.py @@ -994,17 +994,13 @@ async def test_configured_run_kwargs_win_on_conflict( assert run.call_args.kwargs["function_invocation_kwargs"] == {"tenant": "server", "fixed": True, "extra": "x"} - async def test_reserved_keys_are_never_accepted( - self, mock_agent: MagicMock, mock_event_queue: MagicMock, mock_task: Task - ) -> None: - executor = A2AExecutor(mock_agent, accepted_kwargs=["session", "middleware", "tenant"]) - metadata = self._forwarded( - client_kwargs={"session": "x", "middleware": ["y"], "tenant": "acme"}, - ) - - run = await self._execute(executor, metadata, mock_event_queue, mock_task) - - assert run.call_args.kwargs["client_kwargs"] == {"tenant": "acme"} + @mark.parametrize( + "reserved", + ["stream", "messages", "options", "session", "middleware", "client_kwargs", "_internal"], + ) + def test_framework_owned_names_cannot_be_accepted(self, mock_agent: MagicMock, reserved: str) -> None: + with raises(ValueError, match=reserved): + A2AExecutor(mock_agent, accepted_kwargs=["tenant", reserved]) @mark.parametrize( "metadata", From bf5d0e9683c857381f5cffe549ffe3b0aa63fcf6 Mon Sep 17 00:00:00 2001 From: droideronline Date: Tue, 13 Oct 2026 13:36:15 +0530 Subject: [PATCH 4/5] Python: A2AExecutor: authorize function_invocation_kwargs and client_kwargs independently A single accepted_kwargs allowlist let a remote caller inject a tool-context key into client_kwargs, which reaches the chat client. Replace it with accepted_function_invocation_kwargs and accepted_client_kwargs, both empty by default. --- .../a2a/agent_framework_a2a/_a2a_executor.py | 59 +++++++++++-------- .../packages/a2a/tests/test_a2a_executor.py | 49 ++++++++++++--- 2 files changed, 78 insertions(+), 30 deletions(-) diff --git a/python/packages/a2a/agent_framework_a2a/_a2a_executor.py b/python/packages/a2a/agent_framework_a2a/_a2a_executor.py index c0497828112..1958b73ac28 100644 --- a/python/packages/a2a/agent_framework_a2a/_a2a_executor.py +++ b/python/packages/a2a/agent_framework_a2a/_a2a_executor.py @@ -29,7 +29,6 @@ # Key in ``SendMessageRequest.metadata`` that carries forwarded ``function_invocation_kwargs`` and ``client_kwargs``. AGENT_FRAMEWORK_METADATA_KEY = "agent_framework" -_FORWARDED_KWARGS_NAMES = ("function_invocation_kwargs", "client_kwargs") # Names owned by the framework (chat client and tool invocation parameters). A remote caller must never be able to # set them, since a duplicate keyword would fail the task or override framework behavior. _RESERVED_KWARGS = frozenset({ @@ -106,10 +105,13 @@ class A2AExecutor(AgentExecutor): agent: The AI agent to execute. stream: Whether to stream the agent response. Defaults to False. run_kwargs: Additional keyword arguments to pass to the agent's run method. - accepted_kwargs: Names of keys the executor accepts from the caller's request. A client that - forwards ``function_invocation_kwargs`` and ``client_kwargs`` (see ``A2AAgent``) sends them - in the request metadata, and only the keys named here are passed to the agent's run - method. The default (``None``) accepts nothing. + accepted_function_invocation_kwargs: Names of keys the executor accepts from the caller's + forwarded ``function_invocation_kwargs`` (see ``A2AAgent``). A client sends them in the + request metadata, and only the keys named here are passed to the agent's run method. + The default (``None``) accepts nothing. + accepted_client_kwargs: Names of keys the executor accepts from the caller's forwarded + ``client_kwargs``. Authorized independently of ``accepted_function_invocation_kwargs``, + because these values reach the chat client. The default (``None``) accepts nothing. """ def __init__( @@ -117,7 +119,8 @@ def __init__( agent: SupportsAgentRun, stream: bool = False, run_kwargs: Mapping[str, Any] | None = None, - accepted_kwargs: Sequence[str] | None = None, + accepted_function_invocation_kwargs: Sequence[str] | None = None, + accepted_client_kwargs: Sequence[str] | None = None, ): """Initialize the A2AExecutor with the specified agent. @@ -126,16 +129,21 @@ def __init__( stream: Whether to stream the agent response. Defaults to False. run_kwargs: Additional keyword arguments to pass to the agent's run method. Cannot contain 'session' or 'stream' as these are managed by the executor. - accepted_kwargs: Names of keys to accept from the request's forwarded - ``function_invocation_kwargs`` and ``client_kwargs``. Accepted values are passed to the - agent's run method. Values configured in ``run_kwargs`` take precedence on conflict. - Defaults to accepting nothing. Framework-owned names (for example ``stream``, - ``messages``, ``options`` and ``session``) and names starting with an underscore are - rejected. + accepted_function_invocation_kwargs: Names of keys to accept from the request's forwarded + ``function_invocation_kwargs``. Accepted values are passed to the agent's run method. + Values configured in ``run_kwargs`` take precedence on conflict. Defaults to accepting + nothing. + accepted_client_kwargs: Names of keys to accept from the request's forwarded + ``client_kwargs``. Authorized separately from ``accepted_function_invocation_kwargs``, + so opting into a tool-context key does not also let the caller inject it into the + chat client. Defaults to accepting nothing. + + Framework-owned names (for example ``stream``, ``messages``, ``options`` and ``session``) and + names starting with an underscore are rejected in both allowlists. Raises: - ValueError: If run_kwargs contains 'session' or 'stream', or accepted_kwargs contains a - framework-owned name. + ValueError: If run_kwargs contains 'session' or 'stream', or an accepted kwargs allowlist + contains a framework-owned name. """ super().__init__() self._agent: SupportsAgentRun = agent @@ -146,10 +154,17 @@ def __init__( if "stream" in run_kwargs: raise ValueError("run_kwargs cannot contain 'stream' as it is managed by the executor.") self._run_kwargs: Mapping[str, Any] = run_kwargs or {} - reserved = sorted(key for key in accepted_kwargs or () if key in _RESERVED_KWARGS or key.startswith("_")) - if reserved: - raise ValueError(f"accepted_kwargs cannot contain framework-owned names: {', '.join(reserved)}.") - self._accepted_kwargs: frozenset[str] = frozenset(accepted_kwargs or ()) + for allowlist_name, allowlist in ( + ("accepted_function_invocation_kwargs", accepted_function_invocation_kwargs), + ("accepted_client_kwargs", accepted_client_kwargs), + ): + reserved = sorted(key for key in allowlist or () if key in _RESERVED_KWARGS or key.startswith("_")) + if reserved: + raise ValueError(f"{allowlist_name} cannot contain framework-owned names: {', '.join(reserved)}.") + self._accepted_kwargs: dict[str, frozenset[str]] = { + "function_invocation_kwargs": frozenset(accepted_function_invocation_kwargs or ()), + "client_kwargs": frozenset(accepted_client_kwargs or ()), + } @override async def cancel(self, context: RequestContext, event_queue: EventQueue) -> None: @@ -223,18 +238,16 @@ async def execute(self, context: RequestContext, event_queue: EventQueue) -> Non def _build_run_kwargs(self, context: RequestContext) -> dict[str, Any]: """Merge the accepted request kwargs under the configured ``run_kwargs``.""" run_kwargs: dict[str, Any] = dict(self._run_kwargs) - if not self._accepted_kwargs: + if not any(self._accepted_kwargs.values()): return run_kwargs forwarded = context.metadata.get(AGENT_FRAMEWORK_METADATA_KEY) if not isinstance(forwarded, Mapping): return run_kwargs - for name in _FORWARDED_KWARGS_NAMES: + for name, allowed in self._accepted_kwargs.items(): values = cast("Mapping[str, Any]", forwarded).get(name) if not isinstance(values, Mapping): continue - accepted = { - key: value for key, value in cast("Mapping[str, Any]", values).items() if key in self._accepted_kwargs - } + accepted = {key: value for key, value in cast("Mapping[str, Any]", values).items() if key in allowed} if accepted: configured = run_kwargs.get(name) run_kwargs[name] = {**accepted, **(configured or {})} diff --git a/python/packages/a2a/tests/test_a2a_executor.py b/python/packages/a2a/tests/test_a2a_executor.py index fce37fe84de..9af8add0d4e 100644 --- a/python/packages/a2a/tests/test_a2a_executor.py +++ b/python/packages/a2a/tests/test_a2a_executor.py @@ -957,7 +957,9 @@ def _forwarded(**kwargs: Any) -> dict[str, Any]: async def test_accepted_keys_are_passed_to_agent_run( self, mock_agent: MagicMock, mock_event_queue: MagicMock, mock_task: Task ) -> None: - executor = A2AExecutor(mock_agent, accepted_kwargs=["tenant", "trace"]) + executor = A2AExecutor( + mock_agent, accepted_function_invocation_kwargs=["tenant"], accepted_client_kwargs=["trace"] + ) metadata = self._forwarded( function_invocation_kwargs={"tenant": "acme", "other": 1}, client_kwargs={"trace": {"id": "abc"}, "other": 2}, @@ -969,6 +971,34 @@ async def test_accepted_keys_are_passed_to_agent_run( assert kwargs["function_invocation_kwargs"] == {"tenant": "acme"} assert kwargs["client_kwargs"] == {"trace": {"id": "abc"}} + async def test_each_destination_is_authorized_independently( + self, mock_agent: MagicMock, mock_event_queue: MagicMock, mock_task: Task + ) -> None: + executor = A2AExecutor(mock_agent, accepted_function_invocation_kwargs=["tenant"]) + metadata = self._forwarded( + function_invocation_kwargs={"tenant": "acme"}, + client_kwargs={"tenant": "injected"}, + ) + + run = await self._execute(executor, metadata, mock_event_queue, mock_task) + + assert run.call_args.kwargs["function_invocation_kwargs"] == {"tenant": "acme"} + assert "client_kwargs" not in run.call_args.kwargs + + async def test_client_kwargs_require_their_own_opt_in( + self, mock_agent: MagicMock, mock_event_queue: MagicMock, mock_task: Task + ) -> None: + executor = A2AExecutor(mock_agent, accepted_client_kwargs=["trace"]) + metadata = self._forwarded( + function_invocation_kwargs={"trace": "tool"}, + client_kwargs={"trace": "abc"}, + ) + + run = await self._execute(executor, metadata, mock_event_queue, mock_task) + + assert "function_invocation_kwargs" not in run.call_args.kwargs + assert run.call_args.kwargs["client_kwargs"] == {"trace": "abc"} + async def test_nothing_is_accepted_by_default( self, mock_agent: MagicMock, mock_event_queue: MagicMock, mock_task: Task ) -> None: @@ -986,7 +1016,7 @@ async def test_configured_run_kwargs_win_on_conflict( executor = A2AExecutor( mock_agent, run_kwargs={"function_invocation_kwargs": {"tenant": "server", "fixed": True}}, - accepted_kwargs=["tenant", "extra"], + accepted_function_invocation_kwargs=["tenant", "extra"], ) metadata = self._forwarded(function_invocation_kwargs={"tenant": "caller", "extra": "x"}) @@ -998,9 +1028,12 @@ async def test_configured_run_kwargs_win_on_conflict( "reserved", ["stream", "messages", "options", "session", "middleware", "client_kwargs", "_internal"], ) - def test_framework_owned_names_cannot_be_accepted(self, mock_agent: MagicMock, reserved: str) -> None: + @mark.parametrize("allowlist", ["accepted_function_invocation_kwargs", "accepted_client_kwargs"]) + def test_framework_owned_names_cannot_be_accepted( + self, mock_agent: MagicMock, reserved: str, allowlist: str + ) -> None: with raises(ValueError, match=reserved): - A2AExecutor(mock_agent, accepted_kwargs=["tenant", reserved]) + A2AExecutor(mock_agent, **{allowlist: ["tenant", reserved]}) @mark.parametrize( "metadata", @@ -1010,17 +1043,19 @@ def test_framework_owned_names_cannot_be_accepted(self, mock_agent: MagicMock, r async def test_malformed_metadata_is_ignored( self, mock_agent: MagicMock, mock_event_queue: MagicMock, mock_task: Task, metadata: Any ) -> None: - executor = A2AExecutor(mock_agent, accepted_kwargs=["tenant"]) + executor = A2AExecutor( + mock_agent, accepted_function_invocation_kwargs=["tenant"], accepted_client_kwargs=["tenant"] + ) run = await self._execute(executor, metadata, mock_event_queue, mock_task) assert "client_kwargs" not in run.call_args.kwargs assert "function_invocation_kwargs" not in run.call_args.kwargs - async def test_accepted_kwargs_are_passed_when_streaming( + async def test_accepted_function_invocation_kwargs_are_passed_when_streaming( self, mock_agent: MagicMock, mock_event_queue: MagicMock, mock_task: Task ) -> None: - executor = A2AExecutor(mock_agent, stream=True, accepted_kwargs=["tenant"]) + executor = A2AExecutor(mock_agent, stream=True, accepted_function_invocation_kwargs=["tenant"]) stream = MagicMock() stream.with_transform_hook = MagicMock(return_value=stream) stream.get_final_response = AsyncMock() From 7f0cb4e313c19de65cfccf6c02c9c3c147d54a59 Mon Sep 17 00:00:00 2001 From: droideronline Date: Wed, 14 Oct 2026 16:13:38 +0530 Subject: [PATCH 5/5] Python: A2AExecutor: fix typing in reserved name test Pass the allowlists as explicit keyword arguments so pyright, ty and zuban can check the call. --- python/packages/a2a/tests/test_a2a_executor.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/python/packages/a2a/tests/test_a2a_executor.py b/python/packages/a2a/tests/test_a2a_executor.py index 9af8add0d4e..76594d984b7 100644 --- a/python/packages/a2a/tests/test_a2a_executor.py +++ b/python/packages/a2a/tests/test_a2a_executor.py @@ -1028,12 +1028,11 @@ async def test_configured_run_kwargs_win_on_conflict( "reserved", ["stream", "messages", "options", "session", "middleware", "client_kwargs", "_internal"], ) - @mark.parametrize("allowlist", ["accepted_function_invocation_kwargs", "accepted_client_kwargs"]) - def test_framework_owned_names_cannot_be_accepted( - self, mock_agent: MagicMock, reserved: str, allowlist: str - ) -> None: + def test_framework_owned_names_cannot_be_accepted(self, mock_agent: MagicMock, reserved: str) -> None: + with raises(ValueError, match=reserved): + A2AExecutor(mock_agent, accepted_function_invocation_kwargs=["tenant", reserved]) with raises(ValueError, match=reserved): - A2AExecutor(mock_agent, **{allowlist: ["tenant", reserved]}) + A2AExecutor(mock_agent, accepted_client_kwargs=["tenant", reserved]) @mark.parametrize( "metadata",