fix(a2a): add auto_create_session flag to A2aAgentExecutor

Author: nickchecan

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
 

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(

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