From c122214f42358ce408123eaf82b1bbe5d14046ab Mon Sep 17 00:00:00 2001 From: NickChecan Date: Wed, 15 Apr 2026 16:45:26 -0500 Subject: [PATCH 01/11] fix(a2a): respect auto_create_session flag in A2A executor session preparation --- .../adk/a2a/executor/a2a_agent_executor.py | 22 +++------------- .../a2a/executor/a2a_agent_executor_impl.py | 25 ++++++------------- 2 files changed, 11 insertions(+), 36 deletions(-) diff --git a/src/google/adk/a2a/executor/a2a_agent_executor.py b/src/google/adk/a2a/executor/a2a_agent_executor.py index a9b55f526e..7baaa5f011 100644 --- a/src/google/adk/a2a/executor/a2a_agent_executor.py +++ b/src/google/adk/a2a/executor/a2a_agent_executor.py @@ -328,25 +328,11 @@ async def _prepare_session( run_request: AgentRunRequest, runner: Runner, ): - - session_id = run_request.session_id - # create a new session if not exists - user_id = run_request.user_id - session = await runner.session_service.get_session( - app_name=runner.app_name, - user_id=user_id, - session_id=session_id, + session = await runner._get_or_create_session( + user_id=run_request.user_id, + session_id=run_request.session_id, ) - if session is None: - session = await runner.session_service.create_session( - app_name=runner.app_name, - user_id=user_id, - state={}, - session_id=session_id, - ) - # Update run_request with the new session_id - run_request.session_id = session.id - + run_request.session_id = session.id return session def _check_new_version_extension(self, context: RequestContext): diff --git a/src/google/adk/a2a/executor/a2a_agent_executor_impl.py b/src/google/adk/a2a/executor/a2a_agent_executor_impl.py index 320af124df..d04116a931 100644 --- a/src/google/adk/a2a/executor/a2a_agent_executor_impl.py +++ b/src/google/adk/a2a/executor/a2a_agent_executor_impl.py @@ -281,25 +281,14 @@ async def _resolve_session( run_request: AgentRunRequest, runner: Runner, ): - session_id = run_request.session_id - # create a new session if not exists - user_id = run_request.user_id - session = await runner.session_service.get_session( - app_name=runner.app_name, - user_id=user_id, - session_id=session_id, - # Checking existence doesn't require event history. - config=base_session_service.GetSessionConfig(num_recent_events=0), + session = await runner._get_or_create_session( + user_id=run_request.user_id, + session_id=run_request.session_id, + get_session_config=base_session_service.GetSessionConfig( + num_recent_events=0 + ), ) - if session is None: - session = await runner.session_service.create_session( - app_name=runner.app_name, - user_id=user_id, - state={}, - session_id=session_id, - ) - # Update run_request with the new session_id - run_request.session_id = session.id + run_request.session_id = session.id def _get_invocation_metadata( self, executor_context: ExecutorContext From 79c369e8c108287a8e6416bc8a74ae55a5fe4b81 Mon Sep 17 00:00:00 2001 From: NickChecan Date: Wed, 15 Apr 2026 16:55:31 -0500 Subject: [PATCH 02/11] test(a2a): update executor tests to mock runner._get_or_create_session --- .../a2a/executor/test_a2a_agent_executor.py | 23 +++++++++------- .../executor/test_a2a_agent_executor_impl.py | 27 +++++++------------ 2 files changed, 22 insertions(+), 28 deletions(-) diff --git a/tests/unittests/a2a/executor/test_a2a_agent_executor.py b/tests/unittests/a2a/executor/test_a2a_agent_executor.py index 4f44e1363c..df0af8b944 100644 --- a/tests/unittests/a2a/executor/test_a2a_agent_executor.py +++ b/tests/unittests/a2a/executor/test_a2a_agent_executor.py @@ -224,7 +224,7 @@ async def mock_run_async(**kwargs): @pytest.mark.asyncio async def test_prepare_session_new_session(self): - """Test session preparation when session doesn't exist.""" + """Test session preparation delegates to runner._get_or_create_session.""" run_args = AgentRunRequest( user_id="test-user", session_id=None, @@ -232,11 +232,9 @@ async def test_prepare_session_new_session(self): run_config=Mock(spec=RunConfig), ) - # Mock session service - self.mock_runner.session_service.get_session = AsyncMock(return_value=None) mock_session = Mock() mock_session.id = "new-session-id" - self.mock_runner.session_service.create_session = AsyncMock( + self.mock_runner._get_or_create_session = AsyncMock( return_value=mock_session ) @@ -245,10 +243,13 @@ async def test_prepare_session_new_session(self): self.mock_context, run_args, self.mock_runner ) - # Verify session was created + # Verify session was returned and run_request updated assert result == mock_session - assert run_args.session_id is not None - self.mock_runner.session_service.create_session.assert_called_once() + assert run_args.session_id == "new-session-id" + self.mock_runner._get_or_create_session.assert_called_once_with( + user_id="test-user", + session_id=None, + ) @pytest.mark.asyncio async def test_prepare_session_existing_session(self): @@ -260,10 +261,9 @@ async def test_prepare_session_existing_session(self): run_config=Mock(spec=RunConfig), ) - # Mock session service mock_session = Mock() mock_session.id = "existing-session" - self.mock_runner.session_service.get_session = AsyncMock( + self.mock_runner._get_or_create_session = AsyncMock( return_value=mock_session ) @@ -274,7 +274,10 @@ async def test_prepare_session_existing_session(self): # Verify existing session was returned assert result == mock_session - self.mock_runner.session_service.create_session.assert_not_called() + self.mock_runner._get_or_create_session.assert_called_once_with( + user_id="test-user", + session_id="existing-session", + ) def test_constructor_with_callable_runner(self): """Test constructor with callable runner.""" diff --git a/tests/unittests/a2a/executor/test_a2a_agent_executor_impl.py b/tests/unittests/a2a/executor/test_a2a_agent_executor_impl.py index 940b79a0b9..0be889a462 100644 --- a/tests/unittests/a2a/executor/test_a2a_agent_executor_impl.py +++ b/tests/unittests/a2a/executor/test_a2a_agent_executor_impl.py @@ -97,10 +97,10 @@ async def test_execute_success_new_task(self): new_message=Mock(spec=Content), run_config=Mock(spec=RunConfig), ) - # Mock session service + # Mock _get_or_create_session mock_session = Mock() mock_session.id = "test-session" - self.mock_runner.session_service.get_session = AsyncMock( + self.mock_runner._get_or_create_session = AsyncMock( return_value=mock_session ) @@ -200,10 +200,10 @@ async def test_execute_existing_task(self): run_config=Mock(spec=RunConfig), ) - # Mock session service + # Mock _get_or_create_session mock_session = Mock() mock_session.id = "test-session" - self.mock_runner.session_service.get_session = AsyncMock( + self.mock_runner._get_or_create_session = AsyncMock( return_value=mock_session ) @@ -616,7 +616,7 @@ async def test_execute_missing_user_input(self, mock_handle_user_input): ) mock_handle_user_input.return_value = missing_event - self.mock_runner.session_service.get_session = AsyncMock( + self.mock_runner._get_or_create_session = AsyncMock( return_value=Mock(id="test-session") ) self.mock_request_converter.return_value = AgentRunRequest( @@ -638,12 +638,10 @@ async def test_execute_missing_user_input(self, mock_handle_user_input): @pytest.mark.asyncio async def test_resolve_session_creates_new_session(self): - """Test that _resolve_session creates a new session if it doesn't exist.""" - self.mock_runner.session_service.get_session = AsyncMock(return_value=None) - + """Test that _resolve_session delegates to runner._get_or_create_session.""" new_session = Mock() new_session.id = "new-session-id" - self.mock_runner.session_service.create_session = AsyncMock( + self.mock_runner._get_or_create_session = AsyncMock( return_value=new_session ) @@ -656,17 +654,10 @@ async def test_resolve_session_creates_new_session(self): await self.executor._resolve_session(run_request, self.mock_runner) - self.mock_runner.session_service.get_session.assert_called_once_with( - app_name=self.mock_runner.app_name, - user_id="test-user", - session_id="old-session-id", - config=GetSessionConfig(num_recent_events=0, after_timestamp=None), - ) - self.mock_runner.session_service.create_session.assert_called_once_with( - app_name=self.mock_runner.app_name, + self.mock_runner._get_or_create_session.assert_called_once_with( user_id="test-user", - state={}, session_id="old-session-id", + get_session_config=GetSessionConfig(num_recent_events=0, after_timestamp=None), ) assert run_request.session_id == "new-session-id" From c8270d6c1710681e049d9ba13a84ae896dfa9c81 Mon Sep 17 00:00:00 2001 From: NickChecan Date: Wed, 15 Apr 2026 16:59:37 -0500 Subject: [PATCH 03/11] test(a2a): fix integration server to use auto_create_session=True --- tests/unittests/a2a/integration/server.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/unittests/a2a/integration/server.py b/tests/unittests/a2a/integration/server.py index c965a71091..aa0c9d6314 100644 --- a/tests/unittests/a2a/integration/server.py +++ b/tests/unittests/a2a/integration/server.py @@ -44,6 +44,7 @@ def __init__(self, run_async_fn): app_name="FakeApp", agent=agent, session_service=session_service, + auto_create_session=True, ) self.run_async_fn = run_async_fn From 1857816dd80d428315364691b5062193175e28ef Mon Sep 17 00:00:00 2001 From: NickChecan Date: Wed, 15 Apr 2026 17:53:09 -0500 Subject: [PATCH 04/11] fix(a2a): set auto_create_session=True in to_a2a default runner --- src/google/adk/a2a/utils/agent_to_a2a.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/google/adk/a2a/utils/agent_to_a2a.py b/src/google/adk/a2a/utils/agent_to_a2a.py index 3e8ed461e2..d00fd0aa2f 100644 --- a/src/google/adk/a2a/utils/agent_to_a2a.py +++ b/src/google/adk/a2a/utils/agent_to_a2a.py @@ -142,6 +142,7 @@ async def create_runner() -> Runner: session_service=InMemorySessionService(), memory_service=InMemoryMemoryService(), credential_service=InMemoryCredentialService(), + auto_create_session=True, ) # Create A2A components From b98d14f0aa5c9bc2850d9fa257af429c56e3a74d Mon Sep 17 00:00:00 2001 From: nickchecan Date: Fri, 17 Apr 2026 09:21:33 -0500 Subject: [PATCH 05/11] style: apply pyink formatting to a2a executor tests --- tests/unittests/a2a/executor/test_a2a_agent_executor_impl.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/unittests/a2a/executor/test_a2a_agent_executor_impl.py b/tests/unittests/a2a/executor/test_a2a_agent_executor_impl.py index 0be889a462..11a7b709c2 100644 --- a/tests/unittests/a2a/executor/test_a2a_agent_executor_impl.py +++ b/tests/unittests/a2a/executor/test_a2a_agent_executor_impl.py @@ -657,7 +657,9 @@ async def test_resolve_session_creates_new_session(self): self.mock_runner._get_or_create_session.assert_called_once_with( user_id="test-user", session_id="old-session-id", - get_session_config=GetSessionConfig(num_recent_events=0, after_timestamp=None), + get_session_config=GetSessionConfig( + num_recent_events=0, after_timestamp=None + ), ) assert run_request.session_id == "new-session-id" From cbfe37a610ebc1933572ebc237ab2d314fa71beb Mon Sep 17 00:00:00 2001 From: nickchecan Date: Fri, 17 Apr 2026 09:29:26 -0500 Subject: [PATCH 06/11] fix(a2a): add type guards in _resolve_session for mypy compliance --- src/google/adk/a2a/executor/a2a_agent_executor_impl.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/google/adk/a2a/executor/a2a_agent_executor_impl.py b/src/google/adk/a2a/executor/a2a_agent_executor_impl.py index d04116a931..97dec8c58f 100644 --- a/src/google/adk/a2a/executor/a2a_agent_executor_impl.py +++ b/src/google/adk/a2a/executor/a2a_agent_executor_impl.py @@ -281,6 +281,10 @@ async def _resolve_session( run_request: AgentRunRequest, runner: Runner, ): + if not run_request.user_id: + raise ValueError('user_id must be set in AgentRunRequest') + if not run_request.session_id: + raise ValueError('session_id must be set in AgentRunRequest') session = await runner._get_or_create_session( user_id=run_request.user_id, session_id=run_request.session_id, From cc60c4b7c7a8c5abcb9163179ee761bc253a5589 Mon Sep 17 00:00:00 2001 From: nickchecan Date: Fri, 17 Apr 2026 09:34:36 -0500 Subject: [PATCH 07/11] fix(a2a): add mypy type guards and update runner test assertions --- tests/unittests/a2a/utils/test_agent_to_a2a.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/unittests/a2a/utils/test_agent_to_a2a.py b/tests/unittests/a2a/utils/test_agent_to_a2a.py index a9e2458ebd..8dd69b7552 100644 --- a/tests/unittests/a2a/utils/test_agent_to_a2a.py +++ b/tests/unittests/a2a/utils/test_agent_to_a2a.py @@ -327,6 +327,7 @@ async def test_create_runner_function_creates_runner_correctly( session_service=mock_runner_class.call_args[1]["session_service"], memory_service=mock_runner_class.call_args[1]["memory_service"], credential_service=mock_runner_class.call_args[1]["credential_service"], + auto_create_session=True, ) # Verify the services are of the correct types @@ -391,6 +392,7 @@ async def test_create_runner_function_with_agent_without_name( session_service=mock_runner_class.call_args[1]["session_service"], memory_service=mock_runner_class.call_args[1]["memory_service"], credential_service=mock_runner_class.call_args[1]["credential_service"], + auto_create_session=True, ) @patch("google.adk.a2a.utils.agent_to_a2a.A2aAgentExecutor") From 6168774270b92570ee3865ee56d51b57d9dd290f Mon Sep 17 00:00:00 2001 From: nickchecan Date: Fri, 17 Apr 2026 10:05:58 -0500 Subject: [PATCH 08/11] fix(a2a): fix mypy type errors in executor and agent_to_a2a --- src/google/adk/a2a/executor/a2a_agent_executor.py | 2 +- src/google/adk/a2a/executor/a2a_agent_executor_impl.py | 3 ++- src/google/adk/a2a/utils/agent_to_a2a.py | 7 +++++-- 3 files changed, 8 insertions(+), 4 deletions(-) diff --git a/src/google/adk/a2a/executor/a2a_agent_executor.py b/src/google/adk/a2a/executor/a2a_agent_executor.py index 7baaa5f011..32b2cfa9d2 100644 --- a/src/google/adk/a2a/executor/a2a_agent_executor.py +++ b/src/google/adk/a2a/executor/a2a_agent_executor.py @@ -335,7 +335,7 @@ async def _prepare_session( run_request.session_id = session.id return session - def _check_new_version_extension(self, context: RequestContext): + def _check_new_version_extension(self, context: RequestContext) -> bool: """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) diff --git a/src/google/adk/a2a/executor/a2a_agent_executor_impl.py b/src/google/adk/a2a/executor/a2a_agent_executor_impl.py index 97dec8c58f..228f03c606 100644 --- a/src/google/adk/a2a/executor/a2a_agent_executor_impl.py +++ b/src/google/adk/a2a/executor/a2a_agent_executor_impl.py @@ -18,6 +18,7 @@ from datetime import timezone import inspect import logging +from typing import Any from typing import Awaitable from typing import Callable from typing import Optional @@ -296,7 +297,7 @@ async def _resolve_session( def _get_invocation_metadata( self, executor_context: ExecutorContext - ) -> dict[str, str]: + ) -> dict[str, Any]: return { _get_adk_metadata_key('app_name'): executor_context.app_name, _get_adk_metadata_key('user_id'): executor_context.user_id, diff --git a/src/google/adk/a2a/utils/agent_to_a2a.py b/src/google/adk/a2a/utils/agent_to_a2a.py index d00fd0aa2f..d104afaf9b 100644 --- a/src/google/adk/a2a/utils/agent_to_a2a.py +++ b/src/google/adk/a2a/utils/agent_to_a2a.py @@ -14,6 +14,7 @@ from __future__ import annotations +from contextlib import AbstractAsyncContextManager from contextlib import asynccontextmanager import logging from typing import AsyncIterator @@ -85,7 +86,9 @@ def to_a2a( agent_card: Optional[Union[AgentCard, str]] = None, push_config_store: Optional[PushNotificationConfigStore] = None, runner: Optional[Runner] = None, - lifespan: Optional[Callable[[Starlette], AsyncIterator[None]]] = None, + lifespan: Optional[ + Callable[[Starlette], AbstractAsyncContextManager[None]] + ] = None, ) -> Starlette: """Convert an ADK agent to a A2A Starlette application. @@ -171,7 +174,7 @@ async def create_runner() -> Runner: ) # Build the agent card and configure A2A routes - async def setup_a2a(app: Starlette): + async def setup_a2a(app: Starlette) -> None: # Use provided agent card or build one asynchronously if provided_agent_card is not None: final_agent_card = provided_agent_card From 24732c935bba373e4984bedeb80aae5ca33a90d4 Mon Sep 17 00:00:00 2001 From: nickchecan Date: Fri, 17 Apr 2026 10:06:27 -0500 Subject: [PATCH 09/11] fix(memory): add Task type argument to asyncio.Task annotations --- src/google/adk/memory/vertex_ai_memory_bank_service.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/google/adk/memory/vertex_ai_memory_bank_service.py b/src/google/adk/memory/vertex_ai_memory_bank_service.py index 0906e9a6ba..c95fa06d2e 100644 --- a/src/google/adk/memory/vertex_ai_memory_bank_service.py +++ b/src/google/adk/memory/vertex_ai_memory_bank_service.py @@ -41,7 +41,7 @@ # Strong references to fire-and-forget tasks to prevent garbage collection. # See https://docs.python.org/3/library/asyncio-task.html#creating-tasks -_background_tasks: set[asyncio.Task] = set() +_background_tasks: set[asyncio.Task[None]] = set() _GENERATE_MEMORIES_CONFIG_FALLBACK_KEYS = frozenset({ 'disable_consolidation', @@ -565,7 +565,7 @@ def _get_api_client(self) -> vertexai.AsyncClient: return vertexai.Client(project=self._project, location=self._location).aio -def _log_ingest_task_error(task: asyncio.Task) -> None: +def _log_ingest_task_error(task: asyncio.Task[None]) -> None: """Logs errors from fire-and-forget ingest_events tasks.""" if task.cancelled(): return From b8758bf84662c7633b58338e59b000b74feeb0f2 Mon Sep 17 00:00:00 2001 From: NickChecan Date: Fri, 24 Apr 2026 10:01:00 -0300 Subject: [PATCH 10/11] fix(a2a): add auto_create_session flag to A2aAgentExecutor --- .../adk/a2a/executor/a2a_agent_executor.py | 19 ++++++++++++++- .../a2a/executor/a2a_agent_executor_impl.py | 23 +++++++++++++++---- src/google/adk/a2a/utils/agent_to_a2a.py | 1 - 3 files changed, 37 insertions(+), 6 deletions(-) diff --git a/src/google/adk/a2a/executor/a2a_agent_executor.py b/src/google/adk/a2a/executor/a2a_agent_executor.py index 32b2cfa9d2..db0dde0b7a 100644 --- a/src/google/adk/a2a/executor/a2a_agent_executor.py +++ b/src/google/adk/a2a/executor/a2a_agent_executor.py @@ -38,6 +38,7 @@ from google.adk.runners import Runner from typing_extensions import override +from ...errors.session_not_found_error import SessionNotFoundError from ...utils.context_utils import Aclosing from ..agent.interceptors.new_integration_extension import _NEW_A2A_ADK_INTEGRATION_EXTENSION from ..converters.request_converter import AgentRunRequest @@ -75,12 +76,14 @@ def __init__( config: Optional[A2aAgentExecutorConfig] = None, use_legacy: bool = False, force_new_version: bool = False, + auto_create_session: bool = True, ): super().__init__() self._runner = runner self._config = config or A2aAgentExecutorConfig() self._use_legacy = use_legacy self._force_new_version = force_new_version + self._auto_create_session = auto_create_session self._executor_impl = None async def _resolve_runner(self) -> Runner: @@ -141,6 +144,7 @@ async def execute( self._executor_impl = ExecutorImpl( runner=self._runner, config=self._config, + auto_create_session=self._auto_create_session, ) await self._executor_impl.execute(context, event_queue) return @@ -328,10 +332,23 @@ async def _prepare_session( run_request: AgentRunRequest, runner: Runner, ): - session = await runner._get_or_create_session( + session = await runner.session_service.get_session( + app_name=runner.app_name, user_id=run_request.user_id, session_id=run_request.session_id, ) + if not session: + if self._auto_create_session: + session = await runner.session_service.create_session( + app_name=runner.app_name, + user_id=run_request.user_id, + session_id=run_request.session_id, + ) + else: + raise SessionNotFoundError( + f'Session not found: {run_request.session_id}' + ) + # Update run_request with the new session_id run_request.session_id = session.id return session diff --git a/src/google/adk/a2a/executor/a2a_agent_executor_impl.py b/src/google/adk/a2a/executor/a2a_agent_executor_impl.py index 228f03c606..95ba93e1fa 100644 --- a/src/google/adk/a2a/executor/a2a_agent_executor_impl.py +++ b/src/google/adk/a2a/executor/a2a_agent_executor_impl.py @@ -38,6 +38,7 @@ from a2a.types import TextPart from typing_extensions import override +from ...errors.session_not_found_error import SessionNotFoundError from ...runners import Runner from ...sessions import base_session_service from ...utils.context_utils import Aclosing @@ -70,10 +71,12 @@ def __init__( *, runner: Runner | Callable[..., Runner | Awaitable[Runner]], config: Optional[A2aAgentExecutorConfig] = None, + auto_create_session: bool = True, ): super().__init__() self._runner = runner self._config = config or A2aAgentExecutorConfig() + self._auto_create_session = auto_create_session @override async def cancel(self, context: RequestContext, event_queue: EventQueue): @@ -286,13 +289,25 @@ async def _resolve_session( raise ValueError('user_id must be set in AgentRunRequest') if not run_request.session_id: raise ValueError('session_id must be set in AgentRunRequest') - session = await runner._get_or_create_session( + session = await runner.session_service.get_session( + app_name=runner.app_name, user_id=run_request.user_id, session_id=run_request.session_id, - get_session_config=base_session_service.GetSessionConfig( - num_recent_events=0 - ), + # Checking existence doesn't require event history. + config=base_session_service.GetSessionConfig(num_recent_events=0), ) + if not session: + if self._auto_create_session: + session = await runner.session_service.create_session( + app_name=runner.app_name, + user_id=run_request.user_id, + session_id=run_request.session_id, + ) + else: + raise SessionNotFoundError( + f'Session not found: {run_request.session_id}' + ) + # Update run_request with the new session_id run_request.session_id = session.id def _get_invocation_metadata( diff --git a/src/google/adk/a2a/utils/agent_to_a2a.py b/src/google/adk/a2a/utils/agent_to_a2a.py index d104afaf9b..ddeb3129ce 100644 --- a/src/google/adk/a2a/utils/agent_to_a2a.py +++ b/src/google/adk/a2a/utils/agent_to_a2a.py @@ -145,7 +145,6 @@ async def create_runner() -> Runner: session_service=InMemorySessionService(), memory_service=InMemoryMemoryService(), credential_service=InMemoryCredentialService(), - auto_create_session=True, ) # Create A2A components From 3fdbd79ddd5687d6b1dcf4da00fa440f372c5a3c Mon Sep 17 00:00:00 2001 From: NickChecan Date: Fri, 24 Apr 2026 10:12:19 -0300 Subject: [PATCH 11/11] test(a2a): update executor tests to reflect direct session_service calls --- .../a2a/executor/test_a2a_agent_executor.py | 30 ++++++++++++------ .../executor/test_a2a_agent_executor_impl.py | 31 +++++++++++-------- .../unittests/a2a/utils/test_agent_to_a2a.py | 2 -- 3 files changed, 38 insertions(+), 25 deletions(-) diff --git a/tests/unittests/a2a/executor/test_a2a_agent_executor.py b/tests/unittests/a2a/executor/test_a2a_agent_executor.py index df0af8b944..3ced00b21f 100644 --- a/tests/unittests/a2a/executor/test_a2a_agent_executor.py +++ b/tests/unittests/a2a/executor/test_a2a_agent_executor.py @@ -224,17 +224,18 @@ async def mock_run_async(**kwargs): @pytest.mark.asyncio async def test_prepare_session_new_session(self): - """Test session preparation delegates to runner._get_or_create_session.""" + """Test session preparation creates a new session when none is found.""" run_args = AgentRunRequest( user_id="test-user", - session_id=None, + session_id="new-session-id", new_message=Mock(spec=Content), run_config=Mock(spec=RunConfig), ) mock_session = Mock() mock_session.id = "new-session-id" - self.mock_runner._get_or_create_session = AsyncMock( + self.mock_runner.session_service.get_session = AsyncMock(return_value=None) + self.mock_runner.session_service.create_session = AsyncMock( return_value=mock_session ) @@ -243,17 +244,23 @@ async def test_prepare_session_new_session(self): self.mock_context, run_args, self.mock_runner ) - # Verify session was returned and run_request updated + # Verify session was created and run_request updated assert result == mock_session assert run_args.session_id == "new-session-id" - self.mock_runner._get_or_create_session.assert_called_once_with( + self.mock_runner.session_service.get_session.assert_called_once_with( + app_name="test-app", user_id="test-user", - session_id=None, + session_id="new-session-id", + ) + self.mock_runner.session_service.create_session.assert_called_once_with( + app_name="test-app", + user_id="test-user", + session_id="new-session-id", ) @pytest.mark.asyncio async def test_prepare_session_existing_session(self): - """Test session preparation when session exists.""" + """Test session preparation returns existing session without creating one.""" run_args = AgentRunRequest( user_id="test-user", session_id="existing-session", @@ -263,21 +270,24 @@ async def test_prepare_session_existing_session(self): mock_session = Mock() mock_session.id = "existing-session" - self.mock_runner._get_or_create_session = AsyncMock( + self.mock_runner.session_service.get_session = AsyncMock( return_value=mock_session ) + self.mock_runner.session_service.create_session = AsyncMock() # Execute result = await self.executor._prepare_session( self.mock_context, run_args, self.mock_runner ) - # Verify existing session was returned + # Verify existing session was returned without creating a new one assert result == mock_session - self.mock_runner._get_or_create_session.assert_called_once_with( + self.mock_runner.session_service.get_session.assert_called_once_with( + app_name="test-app", user_id="test-user", session_id="existing-session", ) + self.mock_runner.session_service.create_session.assert_not_called() def test_constructor_with_callable_runner(self): """Test constructor with callable runner.""" diff --git a/tests/unittests/a2a/executor/test_a2a_agent_executor_impl.py b/tests/unittests/a2a/executor/test_a2a_agent_executor_impl.py index 11a7b709c2..b0af5df7a2 100644 --- a/tests/unittests/a2a/executor/test_a2a_agent_executor_impl.py +++ b/tests/unittests/a2a/executor/test_a2a_agent_executor_impl.py @@ -97,10 +97,10 @@ async def test_execute_success_new_task(self): new_message=Mock(spec=Content), run_config=Mock(spec=RunConfig), ) - # Mock _get_or_create_session + # Mock session lookup returning existing session mock_session = Mock() mock_session.id = "test-session" - self.mock_runner._get_or_create_session = AsyncMock( + self.mock_runner.session_service.get_session = AsyncMock( return_value=mock_session ) @@ -200,10 +200,10 @@ async def test_execute_existing_task(self): run_config=Mock(spec=RunConfig), ) - # Mock _get_or_create_session + # Mock session lookup returning existing session mock_session = Mock() mock_session.id = "test-session" - self.mock_runner._get_or_create_session = AsyncMock( + self.mock_runner.session_service.get_session = AsyncMock( return_value=mock_session ) @@ -616,7 +616,7 @@ async def test_execute_missing_user_input(self, mock_handle_user_input): ) mock_handle_user_input.return_value = missing_event - self.mock_runner._get_or_create_session = AsyncMock( + self.mock_runner.session_service.get_session = AsyncMock( return_value=Mock(id="test-session") ) self.mock_request_converter.return_value = AgentRunRequest( @@ -638,28 +638,33 @@ async def test_execute_missing_user_input(self, mock_handle_user_input): @pytest.mark.asyncio async def test_resolve_session_creates_new_session(self): - """Test that _resolve_session delegates to runner._get_or_create_session.""" + """Test that _resolve_session creates a session when none is found.""" new_session = Mock() new_session.id = "new-session-id" - self.mock_runner._get_or_create_session = AsyncMock( + self.mock_runner.session_service.get_session = AsyncMock(return_value=None) + self.mock_runner.session_service.create_session = AsyncMock( return_value=new_session ) run_request = AgentRunRequest( user_id="test-user", - session_id="old-session-id", + session_id="new-session-id", new_message=Mock(spec=Content), run_config=Mock(spec=RunConfig), ) await self.executor._resolve_session(run_request, self.mock_runner) - self.mock_runner._get_or_create_session.assert_called_once_with( + self.mock_runner.session_service.get_session.assert_called_once_with( + app_name="test-app", user_id="test-user", - session_id="old-session-id", - get_session_config=GetSessionConfig( - num_recent_events=0, after_timestamp=None - ), + session_id="new-session-id", + config=GetSessionConfig(num_recent_events=0, after_timestamp=None), + ) + self.mock_runner.session_service.create_session.assert_called_once_with( + app_name="test-app", + user_id="test-user", + session_id="new-session-id", ) assert run_request.session_id == "new-session-id" diff --git a/tests/unittests/a2a/utils/test_agent_to_a2a.py b/tests/unittests/a2a/utils/test_agent_to_a2a.py index 8dd69b7552..a9e2458ebd 100644 --- a/tests/unittests/a2a/utils/test_agent_to_a2a.py +++ b/tests/unittests/a2a/utils/test_agent_to_a2a.py @@ -327,7 +327,6 @@ async def test_create_runner_function_creates_runner_correctly( session_service=mock_runner_class.call_args[1]["session_service"], memory_service=mock_runner_class.call_args[1]["memory_service"], credential_service=mock_runner_class.call_args[1]["credential_service"], - auto_create_session=True, ) # Verify the services are of the correct types @@ -392,7 +391,6 @@ async def test_create_runner_function_with_agent_without_name( session_service=mock_runner_class.call_args[1]["session_service"], memory_service=mock_runner_class.call_args[1]["memory_service"], credential_service=mock_runner_class.call_args[1]["credential_service"], - auto_create_session=True, ) @patch("google.adk.a2a.utils.agent_to_a2a.A2aAgentExecutor")